"""
har_iv_chains.py - the option-chain tests of the article, reproducible.

Source. Every chain comes from the public Dolt database
post-no-preference/options, branch master, table option_chain, read through
the DoltHub SQL API:

    https://www.dolthub.com/repositories/post-no-preference/options
    https://www.dolthub.com/api/v1alpha1/post-no-preference/options/master

Columns used: date, act_symbol, expiration, strike, call_put, bid, ask, vol.
The article's data was downloaded on September 6 to 8, 2026. The database
keeps growing and is versioned, so the script prints the commit it reads
from. Each day is cached as a CSV in ./cache, which makes a rerun cheap and
lets a reader keep the exact rows used.

Implied variance for one day (variance_for and svix_for_day):
  1. mid = (bid + ask) / 2 for every contract.
  2. For each expiry, the forward comes from put-call parity: the median of
     strike + (call - put) * exp(r t) over the 8 strikes where call and put
     are closest, with r = 0.04.
  3. Quote filter: keep a contract only if bid > 0, mid >= 0.05 and
     (ask - bid) / mid <= 0.25.
  4. Keep out-of-the-money puts below the forward and calls at or above it.
     An expiry needs at least 8 such strikes.
  5. Interpolate implied volatility linearly in moneyness K / F on a grid
     from 1 / 1.6 to 1.6, price each grid point with Black-76, and integrate:
     variance = 2 / (t F^2) * integral of the OTM price over strike.
  6. Interpolate total variance linearly between the two expiries around
     30 calendar days and annualize. At least 2 priced expiries are needed.

Weekly cross-section (the default run). Wednesdays from January 1, 2020, to
September 4, 2026. Ten symbols were fixed before any was tested (PRESET); four
were added later from the same candidate list (ADDED). A symbol is excluded
before any forecast is fitted when fewer than 120 Wednesdays produce an
implied variance after the quote filter. Realized variance uses daily bars
from yfinance (auto-adjusted), with the same estimator as the article.

Daily series (--daily SPY). Every weekday from January 1, 2019, with a
2,001-point grid. The database repeats the previous session's quotes on some
market holidays, so a day whose shared contracts are at least 98 percent
unchanged from the previous kept day is dropped (is_repeat). The SPY
development and holdout tests in the article were run on this series.

    python har_iv_chains.py              all fourteen symbols, weekly
    python har_iv_chains.py AAPL MSFT    only these, weekly
    python har_iv_chains.py --daily SPY  daily implied series to CSV

Reads only. Never trades.
"""
import json
import math
import os
import sys
import time
import urllib.parse
import urllib.request
from datetime import date, timedelta

import numpy as np
import pandas as pd

API = "https://www.dolthub.com/api/v1alpha1/post-no-preference/options/master"
CACHE = "cache"
OUT = "har_iv_chains.csv"

#--- the 41 large US option markets probed, and the 31 of them the database
#--- carried on September 4, 2026. The fourteen tested come from the 31
PROBED = ["SPY", "DIA", "QQQ", "IWM", "EEM", "EFA", "TLT", "HYG", "XLF", "XLE",
          "XLK", "XLV", "GLD", "SLV", "USO", "VXX",
          "AAPL", "MSFT", "AMZN", "GOOGL", "META", "NVDA", "TSLA", "JPM",
          "BAC", "XOM", "JNJ", "PFE", "WMT", "KO", "DIS", "INTC", "CSCO",
          "T", "VZ", "GE", "BA", "CAT", "MCD", "NKE", "IBM"]
CANDIDATES = ["META", "MSFT", "XLK", "CAT", "XLV", "SPY", "AAPL", "TSLA", "XLE",
              "JPM", "GOOGL", "NVDA", "IBM", "MCD", "AMZN", "CSCO", "INTC",
              "XLF", "DIA", "GE", "WMT", "DIS", "BA", "KO", "JNJ", "VZ", "XOM",
              "BAC", "PFE", "T", "NKE"]
PRESET = ["SPY", "DIA", "XLF", "XLE", "XLK", "AAPL", "MSFT", "JPM", "XOM", "KO"]
ADDED = ["NVDA", "GOOGL", "AMZN", "CAT"]

FIRST, LAST = date(2020, 1, 1), date(2026, 9, 4)
DAILY_FIRST = date(2019, 1, 1)
WEEKDAY = 2                                  # Wednesday

RISK_FREE = 0.04
MAX_SPREAD = 0.25
MIN_PRICE = 0.05
MIN_OTM_STRIKES = 8
PARITY_STRIKES = 8
GRID_LIMIT = 0.60
GRID_WEEKLY = 1501
GRID_DAILY = 2001
TARGET_DAYS = 30
TRADING_DAYS = 252
REPEAT_SHARE = 0.98

HORIZONS = [1, 5]
BASE = ["log_rv1", "log_rv5", "log_rv22"]
FULL = BASE + ["log_iv"]
OOS_FRACTION = 0.5           # begin scoring at the midpoint of each series
MIN_SESSIONS = 120


def query(sql, attempts=3):
    url = API + "?" + urllib.parse.urlencode({"q": sql})
    request = urllib.request.Request(url, headers={"User-Agent": "har_iv_chains/1.0"})
    last = None
    for attempt in range(attempts):
        try:
            with urllib.request.urlopen(request, timeout=120) as response:
                payload = json.loads(response.read().decode("utf-8"))
            if payload.get("query_execution_status") == "Success":
                rows = payload.get("rows", [])
                return pd.DataFrame(rows) if rows else pd.DataFrame()
            last = "server: " + str(payload.get("query_execution_message", ""))
        except Exception as problem:
            last = "network: " + str(problem)
        if attempt + 1 < attempts:
            time.sleep(3 * (attempt + 1))
    raise RuntimeError(last)


def print_version():
    """Print the database commit being read, so a rerun can be dated."""
    try:
        head = query("SELECT commit_hash, date FROM dolt_log LIMIT 1")
        if not head.empty:
            print("Database commit %s, %s\n"
                  % (head["commit_hash"].iloc[0], head["date"].iloc[0]))
            return
    except Exception as problem:
        print("Database version not available: %s" % problem)
    print("Database version not reported; the cache dates the rows used.\n")


def normalize(frame):
    frame = frame.copy()
    for column in ("strike", "bid", "ask", "vol"):
        frame[column] = pd.to_numeric(frame[column], errors="coerce")
    for column in ("date", "expiration", "call_put"):
        frame[column] = frame[column].astype(str)
    return frame


def fetch_day(ticker, day):
    """One session of one symbol, cached; a missing day leaves a marker."""
    os.makedirs(CACHE, exist_ok=True)
    path = os.path.join(CACHE, "%s_%s.csv" % (ticker, day))
    miss = path + ".empty"
    if os.path.exists(miss):
        return None
    if os.path.exists(path):
        try:
            frame = pd.read_csv(path)
        except Exception:
            os.remove(path)
        else:
            if not frame.empty:
                return normalize(frame)
            os.remove(path)
    frame = query("SELECT date, expiration, strike, call_put, bid, ask, vol "
                  "FROM option_chain "
                  "WHERE date = '%s' AND act_symbol = '%s'" % (day, ticker))
    if frame.empty:
        open(miss, "w").close()
        return None
    frame.to_csv(path, index=False)
    return normalize(frame)


def norm_cdf(x):
    return 0.5 * (1.0 + math.erf(x / math.sqrt(2.0)))


def b76(fwd, strike, t, vol, is_call):
    if t <= 0 or vol <= 0:
        return max(fwd - strike, 0.0) if is_call else max(strike - fwd, 0.0)
    d1 = (math.log(fwd / strike) + 0.5 * vol * vol * t) / (vol * math.sqrt(t))
    d2 = d1 - vol * math.sqrt(t)
    if is_call:
        return fwd * norm_cdf(d1) - strike * norm_cdf(d2)
    return strike * norm_cdf(-d2) - fwd * norm_cdf(-d1)


def variance_for(group, t, grid_points=GRID_WEEKLY):
    """Risk-neutral variance of one expiry, steps 2 to 5 of the docstring."""
    pairs = group.pivot_table(index="strike", columns="call_put",
                              values="mid").dropna()
    if len(pairs) < 3 or "Call" not in pairs or "Put" not in pairs:
        return None
    pairs = pairs.assign(gap=(pairs["Call"] - pairs["Put"]).abs())
    near = pairs.nsmallest(PARITY_STRIKES, "gap")
    fwd = float((near.index + (near["Call"] - near["Put"])
                 / math.exp(-RISK_FREE * t)).median())
    if not np.isfinite(fwd) or fwd <= 0:
        return None
    usable = group[(group["bid"] > 0) & (group["mid"] >= MIN_PRICE) &
                   ((group["ask"] - group["bid"]) / group["mid"] <= MAX_SPREAD)]
    puts = usable[(usable["call_put"] == "Put") & (usable["strike"] < fwd)]
    calls = usable[(usable["call_put"] == "Call") & (usable["strike"] >= fwd)]
    otm = pd.concat([puts, calls]).sort_values("strike").drop_duplicates("strike")
    if len(otm) < MIN_OTM_STRIKES:
        return None
    mnes = (otm["strike"] / fwd).to_numpy()
    vols = otm["vol"].to_numpy()
    order = np.argsort(mnes)
    mnes, vols = mnes[order], vols[order]
    lo, hi = 1.0 / (1.0 + GRID_LIMIT), 1.0 + GRID_LIMIT
    grid_m = np.linspace(lo, hi, grid_points)
    grid_v = np.interp(grid_m, mnes, vols)
    grid_k = grid_m * fwd
    price = np.array([b76(fwd, k, t, v, k >= fwd) for k, v in zip(grid_k, grid_v)])
    area = np.trapezoid(price, grid_k) if hasattr(np, "trapezoid") \
        else np.trapz(price, grid_k)
    return 2.0 * area / (t * fwd ** 2)


def svix_for_day(frame, grid_points=GRID_WEEKLY):
    """Thirty-day annualized implied variance for one day, step 6."""
    frame = frame.copy()
    frame["mid"] = (frame["bid"] + frame["ask"]) / 2.0
    asof = pd.to_datetime(frame["date"].iloc[0])
    frame["expiration"] = pd.to_datetime(frame["expiration"])
    frame["days"] = (frame["expiration"] - asof).dt.days
    points = []
    for days, group in frame[frame["days"] > 0].groupby("days"):
        var = variance_for(group, days / 365.0, grid_points)
        if var and var > 0:
            points.append((int(days), var))
    if len(points) < 2:
        return None
    points.sort()
    below = [p for p in points if p[0] <= TARGET_DAYS]
    above = [p for p in points if p[0] >= TARGET_DAYS]
    if not below or not above:
        return min(points, key=lambda p: abs(p[0] - TARGET_DAYS))[1]
    d0, v0 = below[-1]
    d1, v1 = above[0]
    if d0 == d1:
        return v0
    w0, w1 = v0 * d0 / 365.0, v1 * d1 / 365.0
    w = w0 + (TARGET_DAYS - d0) / (d1 - d0) * (w1 - w0)
    return w / (TARGET_DAYS / 365.0)


def is_repeat(now, before):
    """True when a day's shared contracts carry the previous day's quotes."""
    if before is None:
        return False
    keys = ["expiration", "strike", "call_put"]
    both = now.merge(before, on=keys, suffixes=("_a", "_b"))
    if len(both) < 0.5 * min(len(now), len(before)):
        return False
    same = (both["bid_a"].eq(both["bid_b"]) &
            both["ask_a"].eq(both["ask_b"]) &
            both["vol_a"].eq(both["vol_b"]))
    return bool(same.mean() >= REPEAT_SHARE)


def implied_series(ticker, days, grid_points=GRID_WEEKLY, drop_repeats=False,
                   report_every=60):
    rows, absent, repeats, began = [], 0, 0, time.time()
    previous = None
    for i, day in enumerate(days, 1):
        try:
            frame = fetch_day(ticker, day.isoformat())
        except Exception as problem:
            print("      %s failed: %s" % (day, problem))
            absent += 1
            continue
        if frame is None:
            absent += 1
            continue
        if drop_repeats:
            current = frame[["expiration", "strike", "call_put", "bid", "ask", "vol"]]
            if is_repeat(current, previous):
                repeats += 1
                continue
            previous = current
        var = svix_for_day(frame, grid_points)
        if var and var > 0:
            rows.append({"date": pd.Timestamp(day), "iv_var": var})
        if i % report_every == 0 or i == len(days):
            rate = (time.time() - began) / i
            print("      %4d/%d  usable %d  absent %d  repeat %d  ~%d min left"
                  % (i, len(days), len(rows), absent, repeats,
                     (len(days) - i) * rate / 60))
    return pd.DataFrame(rows)


def realized_panel(bars):
    """Overnight gap squared plus Garman-Klass, then trailing and forward means."""
    o, h, l, c = (bars[k] for k in ("Open", "High", "Low", "Close"))
    ln_hl = np.log(h / l)
    ln_co = np.log(c / o)
    daily = (np.log(o / c.shift(1)) ** 2
             + 0.5 * ln_hl ** 2 - (2.0 * np.log(2.0) - 1.0) * ln_co ** 2)
    daily = daily.dropna()
    frame = pd.DataFrame(index=daily.index)
    for n in (1, 5, 22):
        frame["rv_back_%d" % n] = daily.rolling(n).mean() * TRADING_DAYS
    flipped = daily[::-1]
    for hz in HORIZONS:
        #--- mean of sessions t+1..t+hz, strictly after the quote date
        frame["rv_fwd_%d" % hz] = (flipped.rolling(hz).mean()[::-1]
                                   * TRADING_DAYS).shift(-1)
    frame = frame.reset_index()
    return frame.rename(columns={frame.columns[0]: "date"})


def ols(X, y):
    beta, *_ = np.linalg.lstsq(X, y, rcond=None)
    return beta, y - X @ beta


def newey_west(X, resid, lags):
    xtx_inv = np.linalg.pinv(X.T @ X)
    u = X * resid[:, None]
    S = u.T @ u
    for j in range(1, lags + 1):
        w = 1.0 - j / (lags + 1.0)
        G = u[j:].T @ u[:-j]
        S += w * (G + G.T)
    return np.sqrt(np.maximum(np.diag(xtx_inv @ S @ xtx_inv), 0.0))


def walk(frame, cols, target, h, start_at):
    """Expanding window; a training row enters only once its window closed."""
    gap = pd.Timedelta(days=int(math.ceil(h * 7.0 / 5.0)) + 1)
    y = frame[target].to_numpy()
    X = np.column_stack([np.ones(len(frame))] + [frame[c].to_numpy() for c in cols])
    rows = []
    for i in range(len(frame)):
        if i < start_at:
            continue
        t = frame["date"].iloc[i]
        usable = (frame["date"] + gap <= t).to_numpy()
        if usable.sum() < 60:
            continue
        beta, _ = ols(X[usable], y[usable])
        rows.append((t, y[i], float(X[i] @ beta), float(y[usable].mean())))
    return pd.DataFrame(rows, columns=["date", "actual", "pred", "mean"])


def oos_r2(a, p, b):
    sse = float(((a - p) ** 2).sum())
    sst = float(((a - b) ** 2).sum())
    return 1.0 - sse / sst if sst > 0 else float("nan")


def clark_west(a, small, big, lags):
    f = (a - small) ** 2 - (a - big) ** 2 + (small - big) ** 2
    X = np.ones((len(f), 1))
    beta, resid = ols(X, f)
    se = newey_west(X, resid, max(0, lags))
    if se[0] <= 0:
        return float("nan"), float("nan")
    t = beta[0] / se[0]
    return float(t), 0.5 * math.erfc(t / math.sqrt(2.0))


def quote_quality(ticker, days, sample=25):
    """Median relative quoted spread over the last cached days, for the report."""
    spreads = []
    for day in days[-sample:]:
        path = os.path.join(CACHE, "%s_%s.csv" % (ticker, day.isoformat()))
        if not os.path.exists(path):
            continue
        try:
            frame = pd.read_csv(path)
        except Exception:
            continue
        for column in ("bid", "ask"):
            frame[column] = pd.to_numeric(frame[column], errors="coerce")
        mid = (frame["bid"] + frame["ask"]) / 2.0
        keep = mid > 0
        if keep.any():
            spreads.append(float(((frame["ask"] - frame["bid"])[keep]
                                  / mid[keep]).median()))
    return float(np.median(spreads)) if spreads else float("nan")


def test_symbol(ticker, days):
    print("  %s" % ticker)
    iv = implied_series(ticker, days)
    spread = quote_quality(ticker, days)
    print("      median quoted spread %.1f%%" % (100 * spread))
    if len(iv) < MIN_SESSIONS:
        print("      %d usable Wednesdays, fewer than %d: excluded\n"
              % (len(iv), MIN_SESSIONS))
        return [{"symbol": ticker, "horizon": np.nan, "n": len(iv),
                 "n_oos": 0, "r2_har": np.nan, "r2_hariv": np.nan,
                 "gain": np.nan, "cw_t": np.nan, "cw_p": np.nan,
                 "excluded": "quotes too wide, %.0f%% median spread"
                             % (100 * spread)}]

    import yfinance as yf
    raw = yf.download(ticker, start="2019-09-01", end="2026-09-06",
                      auto_adjust=True, progress=False)
    if isinstance(raw.columns, pd.MultiIndex):
        raw.columns = raw.columns.get_level_values(0)
    bars = raw[["Open", "High", "Low", "Close"]].dropna()
    bars.index = pd.to_datetime(bars.index).tz_localize(None).normalize()

    panel = iv.merge(realized_panel(bars), on="date", how="inner")
    for n in (1, 5, 22):
        panel["log_rv%d" % n] = np.log(panel["rv_back_%d" % n].where(
            panel["rv_back_%d" % n] > 0))
    panel["log_iv"] = np.log(panel["iv_var"].where(panel["iv_var"] > 0))
    for hz in HORIZONS:
        panel["log_rv_fwd_%d" % hz] = np.log(panel["rv_fwd_%d" % hz].where(
            panel["rv_fwd_%d" % hz] > 0))

    out = []
    for h in HORIZONS:
        target = "log_rv_fwd_%d" % h
        sub = panel.dropna(subset=[target] + FULL).reset_index(drop=True)
        if len(sub) < MIN_SESSIONS:
            continue
        start_at = int(len(sub) * OOS_FRACTION)
        b = walk(sub, BASE, target, h, start_at)
        f = walk(sub, FULL, target, h, start_at)
        if len(b) < 40:
            continue
        m = b.merge(f, on=["date", "actual", "mean"], suffixes=("_b", "_f"))
        r_b = oos_r2(m["actual"], m["pred_b"], m["mean"])
        r_f = oos_r2(m["actual"], m["pred_f"], m["mean"])
        t, p = clark_west(m["actual"].to_numpy(), m["pred_b"].to_numpy(),
                          m["pred_f"].to_numpy(), max(0, 2 * (h - 1)))
        out.append({"symbol": ticker, "horizon": h, "n": len(sub),
                    "n_oos": len(m), "r2_har": r_b, "r2_hariv": r_f,
                    "gain": r_f - r_b, "cw_t": t, "cw_p": p})
        print("      h=%-2d  n=%3d  HAR %5.1f%%  HAR-IV %5.1f%%  gain %+5.1f  t %5.2f"
              % (h, len(sub), 100 * r_b, 100 * r_f, 100 * (r_f - r_b), t))
    print()
    return out


def run_daily(ticker):
    days = [DAILY_FIRST + timedelta(days=i)
            for i in range((LAST - DAILY_FIRST).days + 1)]
    days = [d for d in days if d.weekday() < 5]
    print("%s daily, %s to %s, %d weekdays\n" % (ticker, days[0], days[-1], len(days)))
    series = implied_series(ticker, days, GRID_DAILY, drop_repeats=True)
    path = "%s_daily_iv.csv" % ticker.lower()
    series.to_csv(path, index=False)
    print("\nWrote %s, %d sessions." % (path, len(series)))
    return 0


def main():
    args = sys.argv[1:]
    print_version()
    if len(args) == 2 and args[0] == "--daily":
        return run_daily(args[1].upper())

    wanted = [s.upper() for s in args] or PRESET + ADDED
    days = [FIRST + timedelta(days=i) for i in range((LAST - FIRST).days + 1)]
    days = [d for d in days if d.weekday() == WEEKDAY]
    print("Candidates (%d): %s" % (len(CANDIDATES), ", ".join(CANDIDATES)))
    print("Fixed in advance: %s" % ", ".join(PRESET))
    print("Added later:      %s\n" % ", ".join(ADDED))
    print("%d symbols, %d Wednesdays each.\n" % (len(wanted), len(days)))

    rows = []
    for ticker in wanted:
        try:
            rows.extend(test_symbol(ticker, days))
        except Exception as problem:
            print("  %s failed: %s\n" % (ticker, problem))
    if not rows:
        print("No results.")
        return 1
    done = pd.DataFrame(rows)
    done.to_csv(OUT, index=False)

    gone = done[done["excluded"].notna()] if "excluded" in done else done.iloc[0:0]
    kept = done[done["excluded"].isna()] if "excluded" in done else done
    print("=" * 72)
    print("  Passed:   %s" % ", ".join(kept["symbol"].unique()))
    print("  Excluded: %s" % ", ".join("%s (%s)" % (r["symbol"], r["excluded"])
                                       for _, r in gone.iterrows()))
    for h in HORIZONS:
        block = kept[kept["horizon"] == h].sort_values("gain", ascending=False)
        if block.empty:
            continue
        print("\n  Horizon %d session(s)\n" % h)
        print("  %-8s %6s %9s %9s %8s %7s %8s"
              % ("symbol", "n", "HAR", "HAR-IV", "gain", "CW t", "p"))
        for _, r in block.iterrows():
            print("  %-8s %6d %8.1f%% %8.1f%% %+7.1f %7.2f %8.4f"
                  % (r["symbol"], r["n"], 100 * r["r2_har"],
                     100 * r["r2_hariv"], 100 * r["gain"], r["cw_t"], r["cw_p"]))
    print("\nWrote %s." % OUT)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
