"""
Daily option-chain snapshots for the S&P 500, collected forward from
2026-10-12. Options-implied stock factors have no free history, so the
record starts with this collector:
  - IV skew (Xing, Zhang & Zhao 2010)
  - the call-put IV spread (Cremers & Weinbaum 2010)
  - option/stock volume (Johnson & So 2012)
  - implied-vol level and changes

    uv run python -m research.lab.fetch_options snapshot            # after the close (systemd timer)
    uv run python -m research.lab.fetch_options snapshot --limit 5 --force --out /tmp/x   # smoke test
    uv run python -m research.lab.fetch_options status

Sources (keys ALPACA_API_KEY / ALPACA_API_SECRET from .env, never logged):
  - stock close and volume: Alpaca SIP daily bars, unadjusted, so they
    match the strikes. The free plan serves bars older than 15 minutes.
  - option quotes: /v1beta1/options/snapshots/{underlying}, feed=indicative.
    This is Alpaca's free derived OPRA feed. It gives quote, last trade,
    today's bar, and Alpaca's own IV and greeks. The real OPRA feed returns
    403 until an OPRA agreement is signed.
  - open interest: /v2/options/contracts (OCC, as of the prior day).
  The filter keeps expiries 5-100 calendar days out and strikes at
  0.7-1.3 x the close. The universe is data/universe_sp500.json plus SPY,
  QQQ and IWM.

When: 16:25 America/New_York, Mon-Fri (deploy/systemd/quant-trading-options.timer).
The Alpaca calendar decides the session. The script records the latest
session that closed at least 10 minutes ago. It skips the run if the next
session has already opened (the quotes would belong to another day) and if
the session's file exists, unless --force.

Files (data/options/us/):
  chains/YYYY/chain_<session>.csv.gz   one row per contract: quote, last trade,
                                        day volume, OI, Alpaca IV/greeks, stock
                                        close/volume, r, implied forward F, and
                                        two IVs from this script's own solver
  summary/summary_<session>.csv        one row per underlying: iv_atm30,
                                        skew_xzz, cp_spread, os_volume,
                                        pc_volume, pc_oi, quote quality
  manifest.jsonl                       one line per run
The IVs are European. iv_fwd is Black-76 on the forward implied by put-call
parity at the three strikes nearest the close, so dividends drop out. iv_spot
is Black-Scholes on the close with no dividends. cp_spread uses iv_spot,
which is what Cremers-Weinbaum measure, so a dividend payer shows a negative
spread of about its yield; adjust for that when testing. American early
exercise is ignored, so the summary uses near-the-money strikes. The raw
chains are kept so a better solver can be run later.

NEEDS EVIDENCE: whether the indicative quotes match the OPRA NBBO closely
enough for skew and spread. Compare a sample against another source before
the first factor test. A cross-sectional test needs about 2-3 years of
snapshots.
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import time
from zoneinfo import ZoneInfo

import numpy as np
import pandas as pd
import requests
from scipy.special import ndtr

from research.lab.fetch_alpaca import _keys

NY = ZoneInfo("America/New_York")
DATA = "https://data.alpaca.markets"
TRADE = "https://paper-api.alpaca.markets"
OUT = os.path.join("data", "options", "us")
UNIVERSE = os.path.join("data", "universe_sp500.json")
RF_PATH = os.path.join("data", "research", "ext", "fred_dtb3.csv")
EXTRA = ["SPY", "QQQ", "IWM"]
DTE_MIN, DTE_MAX = 5, 100
MNY_LO, MNY_HI = 0.7, 1.3
MIN_INTERVAL = 0.32                      # <= ~190 requests a minute (limit 200)


class Api:
    def __init__(self) -> None:
        self.s = requests.Session()
        self.s.headers.update(_keys())
        self.next = 0.0
        self.n = 0

    def get(self, url: str, params: dict) -> dict:
        for attempt in range(6):
            wait = self.next - time.time()
            if wait > 0:
                time.sleep(wait)
            self.next = time.time() + MIN_INTERVAL
            try:
                r = self.s.get(url, params=params, timeout=60)
            except requests.RequestException:
                time.sleep(2 ** attempt)
                continue
            self.n += 1
            if r.status_code == 429 or r.status_code >= 500:
                time.sleep(2 ** attempt)
                continue
            if r.status_code != 200:
                raise RuntimeError(f"HTTP {r.status_code} {r.text[:200]}")
            return r.json()
        raise RuntimeError("retries exhausted")

    def pages(self, url: str, params: dict, key: str):
        params = dict(params)
        while True:
            j = self.get(url, params)
            yield j.get(key) or ({} if key in ("snapshots", "bars") else [])
            tok = j.get("next_page_token")
            if not tok:
                return
            params["page_token"] = tok


# ---------------------------------------------------------------------------
# session and inputs
# ---------------------------------------------------------------------------

def pick_session(api: Api, now: dt.datetime, force: bool) -> dt.date | None:
    today = now.astimezone(NY).date()
    cal = api.get(f"{TRADE}/v2/calendar", {"start": (today - dt.timedelta(days=10)).isoformat(),
                                           "end": (today + dt.timedelta(days=10)).isoformat()})
    days = []
    for c in cal:
        d = dt.date.fromisoformat(c["date"])
        at = lambda hm: dt.datetime.combine(d, dt.time(int(hm[:2]), int(hm[3:5])), NY)  # noqa: E731
        days.append((d, at(c["open"]), at(c["close"])))
    done = [x for x in days if x[2] + dt.timedelta(minutes=10) <= now]
    if not done:
        print("skip: no session in the calendar window has closed")
        return None
    d, _, _ = done[-1]
    nxt = [x for x in days if x[0] > d]
    if nxt and now >= nxt[0][1] and not force:
        print(f"skip: session {d} is over and {nxt[0][0]} has opened; quotes would be from another day")
        return None
    return d


def universe() -> list[str]:
    with open(UNIVERSE) as fh:
        ticks = json.load(fh)["tickers"]
    return sorted({t.replace("-", ".") for t in ticks} | set(EXTRA))


def stock_bars(api: Api, syms: list[str], session: dt.date, now: dt.datetime) -> dict[str, dict]:
    end = min(now - dt.timedelta(minutes=16), dt.datetime.combine(session + dt.timedelta(days=1), dt.time(0), NY))
    out = {}
    for i in range(0, len(syms), 100):
        params = {"symbols": ",".join(syms[i:i + 100]), "timeframe": "1Day", "start": session.isoformat(),
                  "end": end.isoformat(), "feed": "sip", "adjustment": "raw", "limit": 10000}
        for bars in api.pages(f"{DATA}/v2/stocks/bars", params, "bars"):
            for sym, bl in bars.items():
                for b in bl:
                    if b["t"][:10] == session.isoformat():
                        out[sym] = {"close": b["c"], "volume": b["v"]}
    return out


def open_interest(api: Api, syms: list[str], session: dt.date) -> dict[str, dict]:
    out = {}
    lo, hi = session + dt.timedelta(days=DTE_MIN), session + dt.timedelta(days=DTE_MAX)
    for i in range(0, len(syms), 10):
        params = {"underlying_symbols": ",".join(syms[i:i + 10]), "expiration_date_gte": lo.isoformat(),
                  "expiration_date_lte": hi.isoformat(), "limit": 10000}
        try:
            for page in api.pages(f"{TRADE}/v2/options/contracts", params, "option_contracts"):
                for c in page:
                    out[c["symbol"]] = {"oi": float(c["open_interest"]) if c.get("open_interest") else np.nan,
                                        "oi_date": c.get("open_interest_date")}
        except RuntimeError as exc:
            print(f"open interest {syms[i]}..: {exc}", flush=True)
    return out


def chain(api: Api, sym: str, close: float, session: dt.date) -> list[dict]:
    params = {"feed": "indicative", "limit": 1000,
              "expiration_date_gte": (session + dt.timedelta(days=DTE_MIN)).isoformat(),
              "expiration_date_lte": (session + dt.timedelta(days=DTE_MAX)).isoformat(),
              "strike_price_gte": round(MNY_LO * close, 2), "strike_price_lte": round(MNY_HI * close, 2)}
    rows = []
    for snaps in api.pages(f"{DATA}/v1beta1/options/snapshots/{sym}", params, "snapshots"):
        for occ, s in snaps.items():
            q, t, b, g = s.get("latestQuote") or {}, s.get("latestTrade") or {}, s.get("dailyBar") or {}, s.get("greeks") or {}
            today = b.get("t", "")[:10] == session.isoformat()
            rows.append({"underlying": sym, "symbol": occ, "type": occ[-9],
                         "expiry": f"20{occ[-15:-13]}-{occ[-13:-11]}-{occ[-11:-9]}",
                         "strike": int(occ[-8:]) / 1000, "bid": q.get("bp"), "ask": q.get("ap"),
                         "bid_size": q.get("bs"), "ask_size": q.get("as"), "quote_time": q.get("t"),
                         "last": t.get("p"), "last_time": t.get("t"), "day_volume": b.get("v", 0) if today else 0,
                         "alp_iv": s.get("impliedVolatility"), "delta": g.get("delta"), "gamma": g.get("gamma"),
                         "theta": g.get("theta"), "vega": g.get("vega")})
    return rows


def rate(session: dt.date) -> tuple[float, str]:
    try:
        rf = pd.read_csv(RF_PATH, index_col=0, parse_dates=True).iloc[:, 0]
        rf = pd.to_numeric(rf, errors="coerce").dropna().loc[:pd.Timestamp(session)]
        return float(rf.iloc[-1]) / 100, rf.index[-1].date().isoformat()
    except (OSError, IndexError):
        return 0.04, "default"


# ---------------------------------------------------------------------------
# implied volatility and per-underlying summary
# ---------------------------------------------------------------------------

def black76(F, K, T, sig, call, disc):
    sd = sig * np.sqrt(T)
    d1 = (np.log(F / K) + 0.5 * sd ** 2) / sd
    d2 = d1 - sd
    return np.where(call, disc * (F * ndtr(d1) - K * ndtr(d2)), disc * (K * ndtr(-d2) - F * ndtr(-d1)))


def implied_vol(price, F, K, T, call, disc, iters: int = 80) -> np.ndarray:
    lo, hi = np.full(len(price), 1e-4), np.full(len(price), 5.0)
    intrinsic = disc * np.where(call, np.maximum(F - K, 0), np.maximum(K - F, 0))
    ok = np.isfinite(price) & (price > intrinsic + 1e-6) & (price < black76(F, K, T, hi, call, disc))
    for _ in range(iters):
        mid = (lo + hi) / 2
        above = black76(F, K, T, mid, call, disc) > price
        hi, lo = np.where(above, mid, hi), np.where(above, lo, mid)
    return np.where(ok, (lo + hi) / 2, np.nan)


def add_iv(df: pd.DataFrame, r: float) -> pd.DataFrame:
    df["mid"] = (df["bid"] + df["ask"]) / 2
    good = (df["bid"] > 0) & (df["ask"] > df["bid"])
    df.loc[~good, "mid"] = np.nan
    df["T"] = df["dte"] / 365.0
    disc = np.exp(-r * df["T"])
    df["fwd"] = np.nan
    for (_, _), g in df.groupby(["underlying", "expiry"]):
        piv = g.pivot_table(index="strike", columns="type", values="mid")
        if {"C", "P"} <= set(piv.columns):
            piv = piv.dropna()
            if len(piv):
                S = g["stock_close"].iloc[0]
                near = piv.iloc[np.argsort(np.abs(piv.index.to_numpy() - S))[:3]]
                f = near.index.to_numpy() + np.exp(r * g["T"].iloc[0]) * (near["C"] - near["P"]).to_numpy()
                df.loc[g.index, "fwd"] = float(np.median(f))
    nofwd = df["fwd"].isna()
    df.loc[nofwd, "fwd"] = df.loc[nofwd, "stock_close"] * np.exp(r * df.loc[nofwd, "T"])
    call = (df["type"] == "C").to_numpy()
    args = (df["strike"].to_numpy(), df["T"].to_numpy(), call, disc.to_numpy())
    df["iv_fwd"] = implied_vol(df["mid"].to_numpy(), df["fwd"].to_numpy(), *args)
    spot_f = (df["stock_close"] * np.exp(r * df["T"])).to_numpy()
    df["iv_spot"] = implied_vol(df["mid"].to_numpy(), spot_f, *args)
    return df


def summarize(g: pd.DataFrame) -> dict:
    S = g["stock_close"].iloc[0]
    v = g[g["iv_fwd"].notna()]
    out = {"underlying": g["underlying"].iloc[0], "stock_close": S, "stock_volume": g["stock_volume"].iloc[0],
           "n_contracts": len(g), "n_valid": len(v),
           "med_rel_spread": float(((v["ask"] - v["bid"]) / v["mid"]).median()) if len(v) else np.nan}
    atm = {}
    for e, x in v.groupby("expiry"):
        F = x["fwd"].iloc[0]
        k = x["strike"].iloc[np.argmin(np.abs(x["strike"].to_numpy() - F))]
        atm[int(x["dte"].iloc[0])] = float(x.loc[x["strike"] == k, "iv_fwd"].mean())
    ds = sorted(d for d in atm if 7 <= d <= DTE_MAX)
    if ds:
        lo = [d for d in ds if d <= 30]
        hi = [d for d in ds if d >= 30]
        if lo and hi and lo[-1] != hi[0]:          # total variance, linear in time
            a, b = lo[-1], hi[0]
            w = atm[a] ** 2 * a + (atm[b] ** 2 * b - atm[a] ** 2 * a) * (30 - a) / (b - a)
            out["iv_atm30"] = float(np.sqrt(w / 30)) if w > 0 else np.nan
        else:                                       # 30 days exactly, or only one side: nearest
            out["iv_atm30"] = atm[hi[0]] if hi else atm[lo[-1]]
    mid_exp = v[(v["dte"] >= 10) & (v["dte"] <= 60)]
    if len(mid_exp):
        e = mid_exp.loc[(mid_exp["dte"] - 30).abs().idxmin(), "expiry"]
        x = mid_exp[mid_exp["expiry"] == e]
        m = x["strike"] / S
        puts = x[(x["type"] == "P") & (m >= 0.80) & (m <= 0.95)]
        calls = x[(x["type"] == "C") & (m >= 0.95) & (m <= 1.05)]
        if len(puts) and len(calls):
            p = puts.loc[(puts["strike"] / S - 0.95).abs().idxmin(), "iv_fwd"]
            c = calls.loc[(calls["strike"] / S - 1.0).abs().idxmin(), "iv_fwd"]
            out["skew_xzz"] = float(p - c)
        pairs = mid_exp[(mid_exp["strike"] / S).between(0.8, 1.2) & mid_exp["iv_spot"].notna()]
        pv = pairs.pivot_table(index=["expiry", "strike"], columns="type", values=["iv_spot", "oi"])
        if ("iv_spot", "C") in pv.columns and ("iv_spot", "P") in pv.columns:
            pv = pv.dropna(subset=[("iv_spot", "C"), ("iv_spot", "P")])
            if len(pv):
                w = pv[("oi", "C")].fillna(0) + pv[("oi", "P")].fillna(0) if ("oi", "C") in pv and ("oi", "P") in pv \
                    else pd.Series(1.0, index=pv.index)
                w = w if w.sum() > 0 else pd.Series(1.0, index=pv.index)
                out["cp_spread"] = float(((pv[("iv_spot", "C")] - pv[("iv_spot", "P")]) * w).sum() / w.sum())
    cv, pv_ = g.loc[g["type"] == "C", "day_volume"].sum(), g.loc[g["type"] == "P", "day_volume"].sum()
    co, po = g.loc[g["type"] == "C", "oi"].sum(), g.loc[g["type"] == "P", "oi"].sum()
    out.update({"call_volume": cv, "put_volume": pv_, "call_oi": co, "put_oi": po,
                "pc_volume": pv_ / cv if cv else np.nan, "pc_oi": po / co if co else np.nan,
                "os_volume": (cv + pv_) * 100 / out["stock_volume"] if out["stock_volume"] else np.nan})
    return out


# ---------------------------------------------------------------------------
# commands
# ---------------------------------------------------------------------------

def snapshot(args) -> None:
    out = args.out or OUT
    now = dt.datetime.now(NY)
    api = Api()
    session = dt.date.fromisoformat(args.session) if args.session else pick_session(api, now, args.force)
    if session is None:
        return
    path = os.path.join(out, "chains", str(session.year), f"chain_{session}.csv.gz")
    if os.path.exists(path) and not args.force:
        print(f"{path} exists; nothing to do")
        return
    syms = universe()[: args.limit] if args.limit else universe()
    t0 = time.time()
    px = stock_bars(api, syms, session, now)
    oi = open_interest(api, [s for s in syms if s in px], session)
    rows, failed = [], []
    for k, sym in enumerate(syms):
        if sym not in px:
            failed.append([sym, "no stock bar"])
            continue
        try:
            rs = chain(api, sym, px[sym]["close"], session)
        except RuntimeError as exc:
            failed.append([sym, str(exc)[:120]])
            continue
        for r in rs:
            r.update(stock_close=px[sym]["close"], stock_volume=px[sym]["volume"], **oi.get(r["symbol"], {}))
        rows += rs
        if (k + 1) % 50 == 0:
            print(f"  {k + 1}/{len(syms)} underlyings, {len(rows)} contracts, {api.n} requests "
                  f"[{time.time() - t0:.0f}s]", flush=True)
    if not rows:
        raise SystemExit("no contracts collected")
    df = pd.DataFrame(rows)
    for c in ("oi", "oi_date"):
        if c not in df:
            df[c] = np.nan
    df.insert(0, "session", session.isoformat())
    df["dte"] = (pd.to_datetime(df["expiry"]) - pd.Timestamp(session)).dt.days
    r, r_date = rate(session)
    df["r"] = r
    df = add_iv(df, r)
    os.makedirs(os.path.dirname(path), exist_ok=True)
    df.drop(columns=["mid", "T"]).to_csv(path + ".tmp", index=False, compression="gzip")
    os.replace(path + ".tmp", path)
    summ = pd.DataFrame([summarize(g) for _, g in df.groupby("underlying")])
    summ.insert(0, "session", session.isoformat())
    spath = os.path.join(out, "summary", f"summary_{session}.csv")
    os.makedirs(os.path.dirname(spath), exist_ok=True)
    summ.to_csv(spath, index=False)
    rec = {"session": session.isoformat(), "started": dt.datetime.fromtimestamp(t0, NY).isoformat(timespec="seconds"),
           "finished": dt.datetime.now(NY).isoformat(timespec="seconds"), "underlyings": len(syms),
           "with_chain": int(df["underlying"].nunique()), "contracts": len(df),
           "valid_iv": int(df["iv_fwd"].notna().sum()), "requests": api.n, "r": r, "r_date": r_date,
           "feed": "indicative", "failed": failed[:50], "n_failed": len(failed), "forced": bool(args.force)}
    with open(os.path.join(out, "manifest.jsonl"), "a") as fh:
        fh.write(json.dumps(rec) + "\n")
    print(f"session {session}: {rec['with_chain']}/{len(syms)} underlyings, {len(df)} contracts "
          f"({rec['valid_iv']} with IV), {api.n} requests, {time.time() - t0:.0f}s, failed {len(failed)} -> {path}")


def status(args) -> None:
    out = args.out or OUT
    man = os.path.join(out, "manifest.jsonl")
    if not os.path.exists(man):
        print("no snapshots yet")
        return
    recs = [json.loads(x) for x in open(man) if x.strip()]
    print(f"{len(recs)} runs, {len({r['session'] for r in recs})} sessions, first {recs[0]['session']}, "
          f"last {recs[-1]['session']}")
    for r in recs[-10:]:
        print(f"  {r['session']} {r['with_chain']}/{r['underlyings']} underlyings {r['contracts']} contracts "
              f"{r['valid_iv']} IV, failed {r['n_failed']}, {r['started']} -> {r['finished']}")


def main() -> None:
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    s = sub.add_parser("snapshot")
    s.add_argument("--limit", type=int, help="first N underlyings only (testing)")
    s.add_argument("--force", action="store_true", help="record even if late or already recorded")
    s.add_argument("--session", help="label the snapshot with this session date (testing)")
    s.add_argument("--out", help="output directory (default data/options/us)")
    s.set_defaults(fn=snapshot)
    t = sub.add_parser("status")
    t.add_argument("--out")
    t.set_defaults(fn=status)
    args = ap.parse_args()
    args.fn(args)


if __name__ == "__main__":
    main()