"""
Real SPY option-chain utilities: market implied-vol smiles + Heston calibration.

Data: spy_options.csv, a snapshot of the SPY option chain (bid/ask/IV for calls & puts at four
maturities) pulled from Yahoo Finance via yfinance. Because free quotes are unsynchronized, we do NOT
trust the raw call/put parity slope for the discount factor; instead we take the forward as the strike
where C - P changes sign (discount-independent) and assume a flat risk-free rate. Market IVs are computed
from mid prices with the forward Black formula, using out-of-the-money options only (calls above the
forward, puts below) which are the most liquid and informative.
"""
import numpy as np
import pandas as pd
from scipy.stats import norm
from scipy.optimize import brentq, least_squares
import heston1993 as H

R_DEFAULT = 0.043                                                      # ~ current short rate (SPY dividends absorbed in the forward)


def load(path="spy_options.csv"):
    df = pd.read_csv(path)
    return df, float(df["spot"].iloc[0]), str(df["asof"].iloc[0])


def black(F, K, D, T, sig, cp):
    if sig <= 0 or T <= 0:
        return D * max(cp * (F - K), 0.0)
    d1 = (np.log(F / K) + 0.5 * sig ** 2 * T) / (sig * np.sqrt(T)); d2 = d1 - sig * np.sqrt(T)
    return D * cp * (F * norm.cdf(cp * d1) - K * norm.cdf(cp * d2))


def black_iv(price, F, K, D, T, cp):
    if price <= D * max(cp * (F - K), 0.0) + 1e-8:
        return np.nan
    try:
        return brentq(lambda s: black(F, K, D, T, s, cp) - price, 1e-4, 5.0, maxiter=200)
    except Exception:
        return np.nan


def smile(df, dte, r=R_DEFAULT, mny_lo=0.85, mny_hi=1.13):
    """Market implied-vol smile at one maturity. Returns T, discount D, forward F and arrays K, iv, mny."""
    d = df[df["dte"] == dte].copy()
    d["mid"] = 0.5 * (d["bid"] + d["ask"])
    d = d[(d["bid"] > 0) & (d["ask"] > 0) & (d["mid"] > 0.05)]
    d = d[(d["ask"] - d["bid"]) / d["mid"] < 0.5]                      # drop very wide (illiquid) quotes
    C = d[d["type"] == "C"][["strike", "mid"]].rename(columns={"mid": "C"})
    P = d[d["type"] == "P"][["strike", "mid"]].rename(columns={"mid": "P"})
    m = C.merge(P, on="strike").sort_values("strike")
    T = dte / 365.0; D = np.exp(-r * T)
    K = m["strike"].values; cp = (m["C"] - m["P"]).values             # forward = strike where C - P = 0
    j = np.where(np.diff(np.sign(cp)))[0][0]
    F = K[j] + (K[j + 1] - K[j]) * (0 - cp[j]) / (cp[j + 1] - cp[j])
    rows = []
    for _, row in d.iterrows():
        Kk = row["strike"]; s = 1 if row["type"] == "C" else -1
        if (s == 1 and Kk < F) or (s == -1 and Kk > F):               # OTM only
            continue
        iv = black_iv(row["mid"], F, Kk, D, T, s)
        if np.isfinite(iv) and 0.03 < iv < 1.5:
            rows.append((Kk, iv, Kk / F))
    sm = pd.DataFrame(rows, columns=["K", "iv", "mny"]).sort_values("mny")
    sm = sm[(sm["mny"] > mny_lo) & (sm["mny"] < mny_hi)]
    return dict(T=T, D=D, F=F, K=sm["K"].values, iv=sm["iv"].values, mny=sm["mny"].values)


def heston_iv(F, K, D, T, v0, kappa, theta, sigma, rho, nphi=700):
    """Model implied vol: price the option under Heston on the forward measure, then invert to Black IV."""
    price = D * H.heston_call(F, K, 0.0, T, v0, kappa, theta, sigma, rho, nphi=nphi)
    cp = 1 if K > F else -1
    if cp == -1:
        price = price - D * (F - K)                                   # call -> put by parity (for OTM puts)
    return black_iv(price, F, K, D, T, cp)


def calibrate(sm, nphi=700):
    """Least-squares Heston calibration to a SINGLE-maturity smile. One smile identifies rho well but not the
    mean-reversion structure (kappa, theta, v0 are under-determined) -- prefer calibrate_surface below."""
    return calibrate_surface([sm], nphi=nphi, per_mat=None)


def _subsample(sm, per_mat, lo=0.90, hi=1.10):
    """Strikes in [lo,hi] moneyness (where IV inversion is reliable), ~per_mat of them spread evenly."""
    keep = (sm["mny"] >= lo) & (sm["mny"] <= hi)
    K, iv = sm["K"][keep], sm["iv"][keep]
    if per_mat is None or len(K) <= per_mat:
        return K, iv
    idx = np.linspace(0, len(K) - 1, per_mat).round().astype(int)
    return K[idx], iv[idx]


def calibrate_surface(smiles, nphi=600, per_mat=26):
    """Joint Heston calibration to several maturities at once (one (v0,kappa,theta,sigma,rho) for the whole
    surface). The term structure of the skew is what identifies kappa/theta; returns params + per-maturity RMSE."""
    blocks = [(sm["F"], sm["D"], sm["T"]) + _subsample(sm, per_mat) for sm in smiles]
    atm = float(np.interp(1.0, smiles[0]["mny"], smiles[0]["iv"]))

    def errs(p):
        out = []
        for F, D, T, Ks, ivm in blocks:
            e = np.array([heston_iv(F, K, D, T, *p, nphi=nphi) for K in Ks]) - ivm
            out.append(e)
        return np.concatenate(out)

    def resid(x):
        p = (x[0] ** 2, x[2], x[1] ** 2, x[3], np.tanh(x[4]))          # v0, kappa, theta, sigma, rho
        return np.nan_to_num(errs(p), nan=0.0)                          # stray un-invertible strikes -> no penalty

    x0 = [atm, atm, 1.5, 0.8, np.arctanh(-0.65)]
    sol = least_squares(resid, x0, method="lm", max_nfev=4000)
    v0, theta, kappa, sigma, rho = sol.x[0] ** 2, sol.x[1] ** 2, sol.x[2], sol.x[3], np.tanh(sol.x[4])
    per = []
    for F, D, T, Ks, ivm in blocks:
        e = np.array([heston_iv(F, K, D, T, v0, kappa, theta, sigma, rho, nphi=nphi) for K in Ks]) - ivm
        per.append(float(np.sqrt(np.nanmean(e ** 2)) * 100))
    rmse = float(np.sqrt(np.nanmean(errs((v0, kappa, theta, sigma, rho)) ** 2)) * 100)
    return dict(v0=v0, kappa=kappa, theta=theta, sigma=sigma, rho=rho, rmse=rmse, rmse_per=per)