"""
Vol surface model for the Strategy Builder.

Builds a smile-by-expiry model from a real option chain slice (option_chain.py)
and projects a shocked scenario surface (spot/IV-level/skew/term shocks +
optional manual per-cell overrides from the frontend grid).
"""
import math
from typing import Any, Callable, Dict, List, Optional, Tuple

import numpy as np


class Surface:
    """Current-market smile, interpolated per expiry over log-moneyness."""

    def __init__(self, spot: float, expiries: List[Dict[str, Any]]):
        self.spot = spot
        self._tenors: List[Tuple[float, Callable[[float], float]]] = []  # (days, smile_fn)
        for exp in expiries:
            points = _smile_points(exp, spot)
            if not points:
                continue
            self._tenors.append((exp["days_to_expiry"], _build_smile_fn(points)))
        self._tenors.sort(key=lambda x: x[0])

    def iv_at(self, strike: float, days: float) -> float:
        if not self._tenors:
            return 0.20
        moneyness = math.log(max(strike, 1e-6) / max(self.spot, 1e-6))

        if days <= self._tenors[0][0]:
            return self._tenors[0][1](moneyness)
        if days >= self._tenors[-1][0]:
            return self._tenors[-1][1](moneyness)

        for (d0, f0), (d1, f1) in zip(self._tenors, self._tenors[1:]):
            if d0 <= days <= d1:
                iv0, iv1 = f0(moneyness), f1(moneyness)
                w = (days - d0) / max(d1 - d0, 1e-6)
                return iv0 + (iv1 - iv0) * w
        return self._tenors[-1][1](moneyness)


class ScenarioSurface:
    """Shocked surface at the scenario horizon: parametric shifts + manual overrides."""

    def __init__(
        self,
        base: Surface,
        scenario_spot: float,
        iv_level_shift: float,
        skew_tilt: float,
        term_shift: float,
        manual_overrides: Optional[Dict[Tuple[int, float], float]] = None,
    ):
        self.base = base
        self.spot = scenario_spot
        self.iv_level_shift = iv_level_shift
        self.skew_tilt = skew_tilt
        self.term_shift = term_shift
        self.manual_overrides = manual_overrides or {}

    def iv_at(self, strike: float, days: float) -> float:
        override = self._match_override(strike, days)
        if override is not None:
            return override
        base_iv = self.base.iv_at(strike, days)
        moneyness = math.log(max(strike, 1e-6) / max(self.spot, 1e-6))
        shocked = (
            base_iv
            + self.iv_level_shift
            + self.skew_tilt * moneyness
            + self.term_shift * (days / 30.0)
        )
        return max(0.01, shocked)

    def _match_override(self, strike: float, days: float) -> Optional[float]:
        if not self.manual_overrides:
            return None
        strike_pct = strike / max(self.spot, 1e-6)
        best = None
        best_dist = None
        for (o_days, o_pct), iv in self.manual_overrides.items():
            dist = abs(o_days - days) / 30.0 + abs(o_pct - strike_pct)
            if best_dist is None or dist < best_dist:
                best_dist, best = dist, iv
        # Only snap to an override if it's reasonably close to the requested cell
        if best_dist is not None and best_dist < 0.08:
            return best
        return None


def _smile_points(expiry: Dict[str, Any], spot: float) -> List[Tuple[float, float]]:
    """Average call/put IV per strike (filtering stale/zero quotes) -> [(log-moneyness, iv), ...]."""
    by_strike: Dict[float, List[float]] = {}
    for row in expiry.get("calls", []) + expiry.get("puts", []):
        if row["iv"] and row["iv"] > 0.01:
            by_strike.setdefault(row["strike"], []).append(row["iv"])
    if not by_strike:
        return []
    points = [
        (math.log(k / spot), sum(v) / len(v))
        for k, v in sorted(by_strike.items())
    ]
    return points


def _build_smile_fn(points: List[Tuple[float, float]]) -> Callable[[float], float]:
    xs = np.array([p[0] for p in points])
    ys = np.array([p[1] for p in points])

    if len(xs) >= 4:
        from scipy.interpolate import CubicSpline
        spline = CubicSpline(xs, ys, extrapolate=False)

        def fn(x: float) -> float:
            if x <= xs[0]:
                return float(ys[0])
            if x >= xs[-1]:
                return float(ys[-1])
            v = spline(x)
            return float(max(0.01, v))
        return fn

    def fn_linear(x: float) -> float:
        return float(max(0.01, np.interp(x, xs, ys)))
    return fn_linear


def build_surface(chain_slice: Dict[str, Any]) -> Surface:
    return Surface(chain_slice["spot"], chain_slice["expiries"])


def apply_scenario(
    surface: Surface,
    spot_shock_pct: float = 0.0,
    iv_level_shift: float = 0.0,
    skew_tilt: float = 0.0,
    term_shift: float = 0.0,
    manual_grid: Optional[List[Dict[str, Any]]] = None,
) -> ScenarioSurface:
    """
    manual_grid: list of {days_to_expiry, strike_pct, iv} cells edited by the user in the
    frontend grid — takes precedence over the parametric shock at nearby (days, strike_pct).
    """
    scenario_spot = surface.spot * (1 + spot_shock_pct / 100.0)
    overrides: Dict[Tuple[int, float], float] = {}
    for cell in (manual_grid or []):
        if cell.get("iv") is not None:
            overrides[(int(cell["days_to_expiry"]), float(cell["strike_pct"]))] = float(cell["iv"])
    return ScenarioSurface(surface, scenario_spot, iv_level_shift, skew_tilt, term_shift, overrides)