#!/usr/bin/env python3
"""Equity-index VRP on XSP (mini-SPX) options, 2013-2026 — the one real-prior frontier.

For each day: find the ~30d ATM straddle (strike where call~=put = the forward), back out
implied vol from the straddle price, compare to forward realized vol. VRP = IV - RV; a
short-vol seller harvests it. 13y spans 2018/2020/2022 tail events. Tests raw VRP Sharpe,
per-year (incl. crashes), tail/skew, AND the tail-managed version (don't sell when IV rising
— the gate that fixed crypto VRP). XSP = cash-settled European -> no early-exercise distortion.
"""
import glob
import math
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from signal_sweep import sharpe_t  # noqa: E402
import torch  # noqa: E402

DEV = "cuda" if torch.cuda.is_available() else "cpu"
DAY_NS = 86_400 * 10**9


def load_xsp():
    import databento as db
    rows = []
    for p in sorted(glob.glob("data/surfer/xsp/*.dbn")):
        try:
            df = db.DBNStore.from_file(p).to_df().reset_index()
        except Exception:
            continue
        if df.empty or "symbol" not in df.columns:
            continue
        df = df[df["close"] > 0]
        rows.append(df[["ts_event", "symbol", "close"]])
    import pandas as pd
    d = pd.concat(rows, ignore_index=True)
    s = d["symbol"].astype(str).str.replace(" ", "", regex=False)            # "XSP240119C00450000"
    d["right"] = s.str[-9]
    d["strike"] = s.str[-8:].astype(float) / 1000.0
    d["expiry"] = s.str[-15:-9]                                              # YYMMDD
    d["day"] = (d["ts_event"].astype("int64") // DAY_NS)
    return d


def to_epoch_day(yymmdd):
    import datetime
    y = 2000 + int(yymmdd[:2]); mo = int(yymmdd[2:4]); da = int(yymmdd[4:6])
    return (datetime.date(y, mo, da) - datetime.date(1970, 1, 1)).days


def main():
    d = load_xsp()
    print(f"loaded {len(d)} XSP option-days, {d['day'].nunique()} trading days")
    d["exp_day"] = d["expiry"].map(to_epoch_day)
    d["ttm"] = d["exp_day"] - d["day"]
    d = d[(d["ttm"] >= 20) & (d["ttm"] <= 45)]                               # ~30d window

    iv_by_day = {}
    S_by_day = {}
    for day, g in d.groupby("day"):
        # pick the expiry closest to 30d
        exp = g.iloc[(g["ttm"] - 30).abs().argsort()].iloc[0]["exp_day"]
        ge = g[g["exp_day"] == exp]
        calls = ge[ge["right"] == "C"].groupby("strike")["close"].mean()    # dedup AM/PM series
        puts = ge[ge["right"] == "P"].groupby("strike")["close"].mean()
        common = calls.index.intersection(puts.index)
        if len(common) < 3:
            continue
        diff = (calls[common] - puts[common]).abs()
        katm = diff.idxmin()                                                 # ATM: where C~=P
        S = float(katm + calls[katm] - puts[katm])                          # parity forward
        straddle = float(calls[katm] + puts[katm])
        T = float(ge["ttm"].iloc[0]) / 365.0
        if S <= 0 or T <= 0:
            continue
        iv = straddle / (0.8 * S * math.sqrt(T))                            # ATM straddle -> implied vol
        iv_by_day[int(day)] = iv; S_by_day[int(day)] = S

    days = np.array(sorted(S_by_day))
    S = np.array([S_by_day[x] for x in days]); IV = np.array([iv_by_day[x] for x in days])
    r = np.zeros(len(days)); r[1:] = np.log(S[1:] / S[:-1])
    # forward 21-day realized vol (annualized) — what the seller faces
    H = 21
    rv_fwd = np.full(len(days), np.nan)
    for t in range(len(days) - H):
        rv_fwd[t] = np.std(r[t + 1:t + 1 + H]) * math.sqrt(252)
    vrp = IV - rv_fwd                                                       # premium (positive = seller wins)
    # short-vol daily P&L proxy: collect implied variance, pay realized squared return
    iv_lag = np.concatenate([[np.nan], IV[:-1]])
    svol = (iv_lag ** 2) / 252.0 - r ** 2
    year = np.array([1970 + x / 365.25 for x in days]).astype(int)
    T_ = lambda x: torch.tensor(np.asarray(x)[np.isfinite(np.asarray(x))], device=DEV, dtype=torch.float64)

    print(f"\n===== EQUITY-INDEX VRP (XSP, {days.min()}..{days.max()}, {len(days)} days) =====")
    print(f"mean implied vol {np.nanmean(IV):.1%}  mean fwd-realized {np.nanmean(rv_fwd):.1%}  "
          f"mean VRP {np.nanmean(vrp):.1%}  (positive => premium exists)")
    print(f"VRP frac>0 (IV>RV): {np.nanmean(vrp[np.isfinite(vrp)]>0):.2f}")
    print(f"\nshort-vol daily P&L: Sharpe {sharpe_t(T_(svol)):+.2f}  skew {float(((svol[np.isfinite(svol)]-np.nanmean(svol))**3).mean()/np.nanstd(svol)**3):+.2f}  worst-day {np.nanmin(svol)/np.nanstd(svol):+.1f}σ")
    print("per-year short-vol: " + " ".join(f"{y}:{sharpe_t(T_(svol[year==y])):+.1f}" for y in range(2013, 2027) if (year == y).sum() > 60))

    # tail-managed: don't sell when implied vol rising over 5d (the crypto-VRP gate)
    rising = np.zeros(len(days))
    for t in range(6, len(days)):
        rising[t] = 0.0 if IV[t - 1] > IV[t - 6] else 1.0
    svol_tm = rising * svol
    print(f"\ntail-managed (don't sell into rising IV): Sharpe {sharpe_t(T_(svol_tm)):+.2f}  "
          f"skew {float(((svol_tm[np.isfinite(svol_tm)]-np.nanmean(svol_tm))**3).mean()/np.nanstd(svol_tm)**3):+.2f}  worst {np.nanmin(svol_tm)/np.nanstd(svol_tm):+.1f}σ")
    print("per-year tail-managed: " + " ".join(f"{y}:{sharpe_t(T_(svol_tm[year==y])):+.1f}" for y in range(2013, 2027) if (year == y).sum() > 60))
    print("\nVERDICT: VRP frac>0 high + short-vol Sharpe>0 + tail-managed improves skew/recent = real, harvestable equity VRP.")


if __name__ == "__main__":
    main()