"""
har_iv_research.py

Reproduces the four-market forecast comparison from free data. Requires only
numpy, pandas and yfinance, and runs in about a minute.

Cboe publishes a thirty-day implied volatility index for several underlyings,
built from their option chains, free and with years of history. This script
pairs four of those indices with the fund each one is built on and asks, out
of sample, whether the index improves a HAR forecast of realized variance.

What it measures, in order:

  1  how noisy each realized-variance estimator is, and whether they are even
     measuring the same quantity
  2  the share of daily variance that occurs while the market is closed
  3  whether adding implied variance improves a HAR forecast, out of sample,
     on four underlyings across three asset classes
  4  the variance risk premium, as an external check against published work
  5  where the option-implied term helps most
  6  whether the tests survive the overlap between neighbouring forecast
     windows, with twice the Newey-West lags and on non-overlapping windows

Nothing here is fitted in sample. Every forecast comes from an expanding
window that refits at each date and has never seen the observation it is
asked to predict.

    python har_iv_research.py

Reads only. Never trades.
"""

import math

import numpy as np
import pandas as pd

#--- underlying paired with the Cboe index built from its options
PAIRS = [("SPY", "^VIX", "S&P 500"),
         ("GLD", "^GVZ", "gold"),
         ("USO", "^OVX", "crude oil"),
         ("QQQ", "^VXN", "Nasdaq 100")]

START, END = "2015-01-01", "2026-09-06"
TRADING_DAYS = 252
HORIZONS = [1, 5, 22]
LOOKBACKS = [1, 5, 22]
OOS_FRACTION = 0.40          # scoring begins here, everything before is training
MIN_TRAIN = 200
COVID = ("2020-02-15", "2020-04-30")

BASE = ["log_rv1", "log_rv5", "log_rv22"]
FULL = BASE + ["log_iv"]


#--- data ------------------------------------------------------------------

def download(ticker):
    """
    Split-adjusted daily bars for one ticker over the whole sample.

    Assumes yfinance is installed. Returns None on an empty download so the
    caller can report the absence rather than fit on nothing.
    """
    import yfinance as yf
    raw = yf.download(ticker, start=START, end=END,
                      auto_adjust=True, progress=False)
    if raw.empty:
        return None
    if isinstance(raw.columns, pd.MultiIndex):
        raw.columns = raw.columns.get_level_values(0)
    raw.index = pd.to_datetime(raw.index).tz_localize(None).normalize()
    return raw


def daily_estimators(bars):
    """
    Four ways to measure one session's variance from a daily bar.

    A squared close-to-close return is unbiased and extremely noisy, being a
    chi-squared estimate with one degree of freedom. Range estimators use the
    high and the low and are far quieter, but they see only the hours the
    market was open. The fourth adds the overnight move back, and is the one
    the MQL5 engine uses.
    """
    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)
    overnight = np.log(o / c.shift(1)) ** 2

    frame = pd.DataFrame(index=bars.index)
    frame["close"] = np.log(c / c.shift(1)) ** 2
    frame["parkinson"] = ln_hl ** 2 / (4.0 * np.log(2.0))
    frame["garman"] = (0.5 * ln_hl ** 2
                       - (2.0 * np.log(2.0) - 1.0) * ln_co ** 2)
    frame["gk_overnight"] = overnight + frame["garman"]
    frame["overnight"] = overnight
    return frame.dropna()


def build_panel(under, index_ticker):
    """
    Trailing and forward realized variance for one fund, joined to its index.

    Assumes the index is quoted in annualized percent. Forward windows start
    the session after the quote, so no row can see its own outcome.
    """
    bars = download(under)
    vix = download(index_ticker)
    if bars is None or vix is None:
        return None
    daily = daily_estimators(bars[["Open", "High", "Low", "Close"]].dropna())
    v = daily["gk_overnight"]

    frame = pd.DataFrame(index=v.index)
    for n in LOOKBACKS:
        frame["rv_back_%d" % n] = v.rolling(n).mean() * TRADING_DAYS
    flipped = v[::-1]
    for hz in HORIZONS:
        #--- forward window starts the session AFTER the quote, never the same day
        frame["rv_fwd_%d" % hz] = (flipped.rolling(hz).mean()[::-1]
                                   * TRADING_DAYS).shift(-1)
    frame["win_from"] = pd.Series(frame.index, index=frame.index).shift(-1)
    frame = frame.reset_index().rename(columns={"index": "date"})
    frame = frame.rename(columns={frame.columns[0]: "date"})

    iv = vix[["Close"]].rename(columns={"Close": "iv_level"}).reset_index()
    iv = iv.rename(columns={iv.columns[0]: "date"})
    iv["iv_var"] = (iv["iv_level"] / 100.0) ** 2

    panel = frame.merge(iv[["date", "iv_var", "iv_level"]], on="date",
                        how="inner")
    for n in LOOKBACKS:
        panel["log_rv%d" % n] = np.log(panel["rv_back_%d" % n].where(
            panel["rv_back_%d" % n] > 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))
    panel["log_iv"] = np.log(panel["iv_var"].where(panel["iv_var"] > 0))
    return panel, daily


#--- estimation ------------------------------------------------------------

def ols(X, y):
    """Least squares coefficients and residuals. Assumes X carries its own constant."""
    beta, *_ = np.linalg.lstsq(X, y, rcond=None)
    return beta, y - X @ beta


def newey_west(X, resid, lags):
    """
    Standard errors that survive overlapping windows.

    Forecasting h days ahead from every session means neighboring residuals
    share most of their outcome by construction. Ordinary standard errors are
    too small in that situation and every t-statistic looks better than it is.
    """
    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 expanding_forecast(frame, cols, target, h, start_at):
    """
    One forecast per session, each from a model refit on data that closed
    before that session. The embargo is what keeps a training observation
    whose own forward window overlaps the forecast date out of the fit.
    """
    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(start_at, len(frame)):
        t = frame["date"].iloc[i]
        usable = (frame["date"] + gap <= t).to_numpy()
        if usable.sum() < MIN_TRAIN:
            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(actual, pred, benchmark):
    """
    Measured against the expanding historical mean, not against zero.
    Volatility is always positive, so beating zero is meaningless.
    """
    sse = float(((actual - pred) ** 2).sum())
    sst = float(((actual - benchmark) ** 2).sum())
    return 1.0 - sse / sst if sst > 0 else float("nan")


def clark_west(actual, small, big, lags):
    """
    The right test when one model nests the other.

    Under the null that the extra term is useless, the larger model still has
    to estimate its coefficient, and that estimation noise makes it forecast
    slightly worse. A plain test of equal accuracy reads that penalty as
    evidence against the larger model. This correction removes it.
    """
    f = (actual - small) ** 2 - (actual - 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))


#--- reporting -------------------------------------------------------------

def section(title):
    """Print a section heading. Assumes nothing."""
    print()
    print("=" * 72)
    print("  " + title)
    print("=" * 72)
    print()


def report_estimators(daily):
    """Compare the four estimators on SPY. Assumes daily holds all four columns."""
    section("1. REALIZED VARIANCE ESTIMATORS, MEASURED ON SPY")
    reference = daily["close"].mean()
    print("  %-14s %8s %9s %9s %9s"
          % ("estimator", "level", "sd log", "ac lag1", "sees all"))
    print("  " + "-" * 52)
    for name in ("close", "parkinson", "garman", "gk_overnight"):
        s = daily[name]
        ls = np.log(s.where(s > 0)).dropna()
        complete = "yes" if name in ("close", "gk_overnight") else "no"
        print("  %-14s %7.2fx %9.3f %9.3f %9s"
              % (name, s.mean() / reference, ls.std(), ls.autocorr(1), complete))
    print()
    print("  Level is the mean against close to close, which is noisy but")
    print("  correct on average. An estimator built only from one session's")
    print("  open, high, low and close cannot see the overnight move, so a")
    print("  level well below one is a measurement of something smaller than")
    print("  total daily variance, however smooth it looks.")
    print()
    share = daily["overnight"].mean() / reference
    print("  Share of daily variance occurring while the market is closed:")
    print("    %.1f%%" % (100 * share))
    print()


def report_forecasts(results):
    """Print the out-of-sample comparison. Assumes one result per fund and horizon."""
    section("3. DOES IMPLIED VARIANCE IMPROVE THE FORECAST")
    print("  Out-of-sample R-squared against the expanding historical mean.")
    print("  HAR uses price history alone. HAR-IV adds one term.\n")
    print("  %-16s %4s %8s %9s %9s %8s %9s"
          % ("underlying", "h", "n", "HAR", "HAR-IV", "gain", "CW p"))
    print("  " + "-" * 66)
    for r in results:
        print("  %-16s %4d %8d %8.1f%% %8.1f%% %+7.1f %9.4f"
              % ("%s / %s" % (r["under"], r["index"]), r["h"], r["n_oos"],
                 100 * r["r2_har"], 100 * r["r2_hariv"],
                 100 * r["gain"], r["cw_p"]))
    print()
    gains = [r["gain"] for r in results]
    print("  Positive on %d of %d tests. Smallest gain %+.1f points."
          % (sum(1 for g in gains if g > 0), len(gains), 100 * min(gains)))
    print()


def report_overlap(results):
    """Print the Clark-West test under two harder treatments of overlap."""
    section("6. DO THE TESTS SURVIVE THE OVERLAP")
    print("  Forecasts are made every session, so for h > 1 neighbouring")
    print("  windows share sessions. The main test uses Newey-West errors over")
    print("  2(h-1) lags. Here it is repeated with 4(h-1) lags, and on every")
    print("  h-th forecast only, where no correction is needed.\n")
    print("  %-16s %4s %11s %11s %6s %13s"
          % ("underlying", "h", "p, 2(h-1)", "p, 4(h-1)", "n", "p, no overlap"))
    print("  " + "-" * 66)
    for r in results:
        print("  %-16s %4d %11.6f %11.6f %6d %13.6f"
              % ("%s / %s" % (r["under"], r["index"]), r["h"], r["cw_p"],
                 r["cw_p_nw4"], r["n_nonoverlap"], r["cw_p_nonoverlap"]))
    print()


def report_premium(panels):
    """Print how often implied variance exceeded the next 22 sessions' realized."""
    section("4. VARIANCE RISK PREMIUM, AN EXTERNAL CHECK")
    print("  Implied variance should exceed subsequent realized variance most")
    print("  of the time. Published work puts the figure near 80 percent.")
    print("  Nothing here was tuned to reproduce it.\n")
    print("  %-16s %8s %12s %14s"
          % ("underlying", "n", "share iv>rv", "median gap, pts"))
    print("  " + "-" * 54)
    for under, panel in panels.items():
        sub = panel[["iv_var", "rv_fwd_22", "iv_level"]].dropna()
        share = float((sub["iv_var"] > sub["rv_fwd_22"]).mean())
        gap = float((sub["iv_level"]
                     - 100 * np.sqrt(sub["rv_fwd_22"])).median())
        print("  %-16s %8d %11.1f%% %14.2f" % (under, len(sub),
                                               100 * share, gap))
    print()


def report_where(results):
    """Relate the gain from the implied term to how well HAR does alone."""
    section("5. WHERE THE OPTION-IMPLIED TERM HELPS MOST")
    block = [r for r in results if r["h"] == 5]
    if len(block) < 3:
        return
    x = np.array([r["r2_har"] for r in block])
    y = np.array([r["gain"] for r in block])
    r = float(np.corrcoef(x, y)[0, 1])
    print("  %-16s %10s %10s %9s" % ("underlying", "HAR", "HAR-IV", "gain"))
    print("  " + "-" * 48)
    for b in sorted(block, key=lambda z: z["r2_har"]):
        print("  %-16s %9.1f%% %9.1f%% %+8.1f"
              % (b["under"], 100 * b["r2_har"], 100 * b["r2_hariv"],
                 100 * b["gain"]))
    print()
    print("  Correlation between how well HAR does alone and how much the")
    print("  implied term adds: %+.3f" % r)
    print()
    print("  Four points cannot establish a relationship. The ten-equity")
    print("  study measures the same thing over ten points.")
    print()


#--- main ------------------------------------------------------------------

def main():
    """
    Run the four funds, audit for lookahead, print every section, save a CSV.

    Assumes network access to yfinance. Stops with status 1 if any forward
    window starts on or before its own quote date.
    """
    print("HAR-IV research script. Volatility-index route.")
    print("Sample %s to %s.\n" % (START, END))

    results, panels, spy_daily = [], {}, None
    for under, index_ticker, label in PAIRS:
        print("  downloading %s and %s (%s) ..." % (under, index_ticker, label))
        try:
            built = build_panel(under, index_ticker)
        except Exception as problem:
            print("    failed: %s" % problem)
            continue
        if built is None:
            print("    no data")
            continue
        panel, daily = built
        if under == "SPY":
            spy_daily = daily
        panels[under] = panel

        #--- lookahead audit, printed rather than assumed
        leaks = int((panel["win_from"] <= panel["date"]).sum())
        if leaks:
            print("    LOOKAHEAD: %d rows. Stopping." % leaks)
            return 1

        for h in HORIZONS:
            target = "log_rv_fwd_%d" % h
            sub = panel.dropna(subset=[target] + FULL).reset_index(drop=True)
            start_at = int(len(sub) * OOS_FRACTION)
            base = expanding_forecast(sub, BASE, target, h, start_at)
            full = expanding_forecast(sub, FULL, target, h, start_at)
            if len(base) < 100:
                continue
            m = base.merge(full, 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)))
            #--- the same test twice more: with twice the lags, and on every
            #--- h-th forecast only, where no two windows share a session
            t4, p4 = clark_west(m["actual"].to_numpy(), m["pred_b"].to_numpy(),
                                m["pred_f"].to_numpy(), max(0, 4 * (h - 1)))
            thin = m.iloc[::h]
            tn, pn = clark_west(thin["actual"].to_numpy(),
                                thin["pred_b"].to_numpy(),
                                thin["pred_f"].to_numpy(), 0)
            results.append({"under": under, "index": index_ticker.lstrip("^"),
                            "h": 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,
                            "cw_t_nw4": t4, "cw_p_nw4": p4,
                            "n_nonoverlap": len(thin), "cw_t_nonoverlap": tn,
                            "cw_p_nonoverlap": pn})
    if not results:
        print("\nNothing ran.\n")
        return 1

    if spy_daily is not None:
        report_estimators(spy_daily)
    section("2. LOOKAHEAD AUDIT")
    print("  Zero rows where a forward window starts on or before its quote")
    print("  date, on every panel built above.")
    print()
    report_forecasts(results)
    report_premium(panels)
    report_where(results)
    report_overlap(results)

    pd.DataFrame(results).to_csv("har_iv_results.csv", index=False)
    print("Wrote har_iv_results.csv.")
    return 0


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