#!/usr/bin/env python3
# MIT License
#
# Copyright (c) 2026 Options Lab (ashares options learning site)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
"""Can a time-series model beat "implied vol minus a premium"? Reproduce the headline numbers.

For learning only. Past results do not guarantee future ones; not investment advice.

What it does
------------
1. Downloads public data with AKShare: daily closes of six A-share ETFs and
   their tracking indices (Sina), and each ETF option's QVIX, a VIX-style
   30-day implied-vol index.
2. Builds one sample per ETF roughly every 30 days from 2015-05-11: the
   forecast date t, QVIX at t, and the target, the realised vol of the daily
   returns in (t, t + 30 calendar days].
3. Scores forecasts of that target:
     * trailing realised vol (mean of the 30 / 60 / 90-day RV),
     * raw QVIX,
     * IV - premium: QVIX minus the mean premium (IV - realised) of all
       samples whose 30-day window had already closed (out of sample),
     * a + b*IV, fitted the same way (a control),
     * GARCH(1,1), refitted at every date on all returns up to it,
     * HAR-RV (daily / weekly / monthly squared returns), trained only on
       windows that had already closed,
     * and the two models with their past bias removed.
   RMSE in vol points on the "common sample" (every forecast present), and a
   Diebold-Mariano t against IV - premium (t > 0 = IV - premium is better).
4. Calibration: a log-normal range around IV - premium whose width is the
   spread of its own past log errors. Share of outcomes inside the 50 / 80 /
   90 % ranges, from July 2024 on and over the full history.
5. Optional (--with-fm): Google TimesFM 2.5, zero-shot, on the daily squared
   return series, against IV - premium from July 2024 on.

Run
---
    pip install -r requirements.txt
    python vol_forecast.py                     # download, then print the numbers (~10 s after the download)
    python vol_forecast.py --cache-dir data    # keep the download as CSV and reuse it next time
    python vol_forecast.py --end none          # use every month available, not just the published run
    python vol_forecast.py --with-fm           # + TimesFM 2.5 (needs the extras below)

TimesFM extras (about 1 GB with PyTorch; CPU is enough, ~1 minute):
    pip install "timesfm[torch]>=3.0.2" "torch>=2.4"
The weights (google/timesfm-2.5-200m-pytorch) download from Hugging Face on
first use; where huggingface.co is unreachable set HF_ENDPOINT to a mirror.

Python 3.10+. The numbers move a little when the data vendor revises old rows
or when `arch` changes its optimiser; see the tolerance on the web page.
"""

from __future__ import annotations

import argparse
import bisect
import math
import statistics
import warnings
from datetime import date, timedelta
from pathlib import Path

import numpy as np
import pandas as pd

# ---------------------------------------------------------------------------
# Data
# ---------------------------------------------------------------------------

# code: (name, QVIX series, AKShare QVIX function, Sina ETF symbol, tracking index, Sina index symbol)
UNDERLYINGS = {
    "510050": ("50ETF", "50etf", "index_option_50etf_qvix", "sh510050", "000016", "sh000016"),
    "510300": ("300ETF", "300etf", "index_option_300etf_qvix", "sh510300", "000300", "sh000300"),
    "510500": ("500ETF", "500etf", "index_option_500etf_qvix", "sh510500", "000905", "sh000905"),
    "159915": ("ChiNext ETF", "cyb", "index_option_cyb_qvix", "sz159915", "399006", "sz399006"),
    "588000": ("STAR 50 ETF", "kcb", "index_option_kcb_qvix", "sh588000", "000688", "sh000688"),
    "159901": ("SZSE 100 ETF", "100etf", "index_option_100etf_qvix", "sz159901", "399330", "sz399330"),
}
CODES = tuple(UNDERLYINGS)


def fetch(cache_dir: str | Path | None = None) -> dict[str, pd.DataFrame]:
    """Download (or reload) the three tables.

    Returns long frames: ``qvix`` (series, date, close), ``etf`` and ``idx``
    (code, date, close). ``close`` is in index points for QVIX (18.5 = 18.5 %).
    """
    names = ("qvix", "etf", "idx")
    if cache_dir is not None:
        paths = {n: Path(cache_dir) / f"{n}.csv" for n in names}
        if all(p.exists() for p in paths.values()):
            return {n: pd.read_csv(p, parse_dates=["date"], dtype={"code": str, "series": str})
                    for n, p in paths.items()}

    import akshare as ak  # imported here so compute() works without it

    qv, etf, idx = [], [], []
    for code, (_, series, qfunc, etf_sym, ix_code, ix_sym) in UNDERLYINGS.items():
        q = getattr(ak, qfunc)().rename(columns=str.lower)[["date", "close"]]
        qv.append(q.assign(series=series))
        e = ak.fund_etf_hist_sina(symbol=etf_sym)[["date", "close"]]
        etf.append(e.assign(code=code))
        i = ak.stock_zh_index_daily(symbol=ix_sym)[["date", "close"]]
        idx.append(i.assign(code=ix_code))
    out = {
        "qvix": pd.concat(qv, ignore_index=True),
        "etf": pd.concat(etf, ignore_index=True),
        "idx": pd.concat(idx, ignore_index=True),
    }
    for df in out.values():
        df["date"] = pd.to_datetime(df["date"])
    if cache_dir is not None:
        Path(cache_dir).mkdir(parents=True, exist_ok=True)
        for n, df in out.items():
            df.to_csv(Path(cache_dir) / f"{n}.csv", index=False)
    return out


def _series(df: pd.DataFrame, col: str, key: str) -> pd.Series:
    s = df[df[col] == key].set_index("date")["close"].astype(float).sort_index()
    return s[~s.index.duplicated(keep="last")]


def qvix_series(frames: dict, code: str) -> pd.Series:
    """QVIX as a decimal (0.185), zero rows dropped."""
    s = _series(frames["qvix"], "series", UNDERLYINGS[code][1])
    return s[s > 0] / 100.0


def corporate_action_dates(px: pd.Series, ix: pd.Series) -> set[date]:
    """Days whose ETF return is a split / consolidation, not a price move.

    The ETF moved > 10 % while its index stayed calm (< 3 %), the gap is > 5 %,
    and it is not a +/-10 % limit move. These returns are left out of every
    realised vol (ordinary ex-dividend drops, < 3 %, are kept).
    """
    df = pd.DataFrame({"px": px, "ix": ix}).dropna()
    r_px, r_ix = np.log(df["px"]).diff(), np.log(df["ix"]).diff()
    at_limit = np.zeros(len(df), dtype=bool)
    for lim in (math.log(0.9), math.log(1.1)):
        at_limit |= (r_px - lim).abs().to_numpy() < 0.001
    hit = ((r_px - r_ix).abs() > 0.05) & (r_ix.abs() < 0.03) & (r_px.abs() > 0.10) & ~at_limit
    return {ts.date() for ts in df.index[hit.fillna(False).to_numpy()]}


# ---------------------------------------------------------------------------
# Realised vol and the samples
# ---------------------------------------------------------------------------

MIN_RETURNS = 10
GRID_START = date(2015, 5, 11)  # 50ETF QVIX starts 2015-02-09, + 90 days of history
STEP_DAYS = 30                   # forecast window, calendar days
OOS_MIN = 24                     # closed samples before an out-of-sample fit starts
FLOOR = 0.01


def realized_vol(points: list[tuple[date, float]], as_of: date, lookback_days: int, *,
                 skip_gaps: bool, drops: set[date]) -> float | None:
    """Annualised close-to-close vol over [as_of - lookback, as_of].

    Population sd of the log returns (divide by n) x sqrt(252). ``skip_gaps``
    drops returns spanning more than 5 calendar days (a long holiday), the rule
    of the trailing forecast; the target keeps every trading-day return.
    """
    start = as_of - timedelta(days=lookback_days)
    window = [(d, p) for d, p in points if start <= d <= as_of and p > 0]
    rets = []
    for (d0, p0), (d1, p1) in zip(window, window[1:]):
        if (skip_gaps and (d1 - d0).days > 5) or d1 in drops:
            continue
        rets.append(math.log(p1 / p0))
    if len(rets) < MIN_RETURNS:
        return None
    m = sum(rets) / len(rets)
    return math.sqrt(sum((r - m) ** 2 for r in rets) / len(rets) * 252)


def origin_grid(calendar: list[date]) -> list[date]:
    """First trading day >= 2015-05-11, then the first trading day >= previous + 30 days."""
    i = bisect.bisect_left(calendar, GRID_START)
    out = []
    while i < len(calendar):
        out.append(calendar[i])
        i = bisect.bisect_left(calendar, calendar[i] + timedelta(days=STEP_DAYS))
    return out


def build_samples(frames: dict) -> pd.DataFrame:
    calendar = [ts.date() for ts in _series(frames["etf"], "code", "510050").index]
    grid = origin_grid(calendar)
    rows = []
    for code in CODES:
        q = qvix_series(frames, code)
        q_map = {ts.date(): v for ts, v in q.items()}
        px = _series(frames["etf"], "code", code)
        drops = corporate_action_dates(px, _series(frames["idx"], "code", UNDERLYINGS[code][4]))
        pts = [(ts.date(), float(v)) for ts, v in px.items()]
        dates = [d for d, _ in pts]
        for t in grid:
            if t not in q_map or t + timedelta(days=STEP_DAYS) > dates[-1]:
                continue
            back = pts[bisect.bisect_left(dates, t - timedelta(days=95)): bisect.bisect_right(dates, t)]
            fwd = pts[bisect.bisect_left(dates, t): bisect.bisect_right(dates, t + timedelta(days=STEP_DAYS))]
            trailing = [realized_vol(back, t, w, skip_gaps=True, drops=drops) for w in (30, 60, 90)]
            trailing = [v for v in trailing if v is not None]
            rv_fwd = realized_vol(fwd, t + timedelta(days=STEP_DAYS), STEP_DAYS, skip_gaps=False, drops=drops)
            if not trailing or rv_fwd is None:
                continue
            rows.append(dict(underlying=code, origin=pd.Timestamp(t), iv=q_map[t],
                             fa=max(sum(trailing) / len(trailing), FLOOR), rv_fwd=rv_fwd))
    return pd.DataFrame(rows).sort_values(["origin", "underlying"]).reset_index(drop=True)


def closed_before(df: pd.DataFrame, t: pd.Timestamp) -> np.ndarray:
    """Samples whose 30-day window has closed by t: the only ones a forecast at t may learn from."""
    return (df["origin"] + pd.Timedelta(days=STEP_DAYS) <= t).to_numpy()


def add_iv_forecasts(df: pd.DataFrame) -> pd.DataFrame:
    """IV - premium and a + b*IV, both fitted on closed samples only (pooled over the ETFs)."""
    df = df.copy()
    prem = np.full(len(df), np.nan)
    lin = np.full(len(df), np.nan)
    iv, y = df["iv"].to_numpy(), df["rv_fwd"].to_numpy()
    for t, idx in df.groupby("origin").groups.items():
        ok = closed_before(df, t)
        if ok.sum() < OOS_MIN:
            continue
        pos = df.index.get_indexer(idx)
        prem[pos] = np.mean(iv[ok] - y[ok])
        b, a = np.polyfit(iv[ok], y[ok], 1)
        lin[pos] = a + b * iv[pos]
    df["premium"] = prem
    df["iv_minus_premium"] = np.maximum(iv - prem, FLOOR)
    df["a_b_iv"] = np.maximum(lin, FLOOR)
    return df


# ---------------------------------------------------------------------------
# GARCH(1,1) and HAR-RV
# ---------------------------------------------------------------------------


def daily_returns(px: pd.Series, drops: set[date]) -> pd.Series:
    """Log returns between consecutive closes (holidays kept), corporate actions removed."""
    r = np.log(px[px > 0]).diff().iloc[1:]
    return r[[d.date() not in drops for d in r.index]]


def garch11_vol(r: np.ndarray, h: int) -> float:
    """Fit GARCH(1,1) (normal, constant mean) on r and forecast the annualised vol of the next h days."""
    from arch import arch_model

    am = arch_model(np.asarray(r) * 100.0, mean="Constant", vol="GARCH", p=1, o=0, q=1,
                    dist="normal", rescale=False)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        res = am.fit(disp="off", show_warning=False)
        var_path = res.forecast(horizon=h, reindex=False).variance.to_numpy()[-1] / 1e4
    return math.sqrt(252 * var_path.mean())


def har_frame(r: pd.Series) -> pd.DataFrame:
    """Per return date d: the HAR inputs from returns <= d, and the target y over (d, d + 30 days].

    rv_d / rv_w / rv_m = 252 x mean(r^2) over the last 1 / 5 / 22 returns.
    y = 252 x mean(r^2) of the returns in the next 30 calendar days; it becomes
    known on y_end = d + 30 days, so a model fitted at t only uses y_end <= t.
    """
    s = r ** 2
    out = pd.DataFrame(index=r.index)
    for col, k in (("rv_d", 1), ("rv_w", 5), ("rv_m", 22)):
        out[col] = 252 * s.rolling(k).mean()
    csum = np.concatenate([[0.0], np.cumsum(s.to_numpy())])
    ends = r.index + pd.Timedelta(days=STEP_DAYS)
    lo = np.arange(1, len(r) + 1)
    hi = r.index.searchsorted(ends, side="right")
    n = hi - lo
    with np.errstate(invalid="ignore", divide="ignore"):
        y = 252 * (csum[hi] - csum[lo]) / n
    out["y"] = np.where((n > 0) & (ends <= r.index[-1]), y, np.nan)
    out["y_end"] = ends
    return out


def har_vol(har: pd.DataFrame, t: pd.Timestamp, min_rows: int = 250) -> float:
    cols = ["rv_d", "rv_w", "rv_m"]
    if t not in har.index or har.loc[t, cols].isna().any():
        return np.nan
    train = har[har["y_end"] <= t].dropna(subset=[*cols, "y"])
    if len(train) < min_rows:
        return np.nan
    X = np.column_stack([np.ones(len(train)), train[cols].to_numpy()])
    coef = np.linalg.lstsq(X, train["y"].to_numpy(), rcond=None)[0]
    yhat = coef[0] + coef[1:] @ har.loc[t, cols].to_numpy(dtype=float)
    return math.sqrt(max(yhat, train["y"].min()))


def add_models(df: pd.DataFrame, frames: dict, garch: bool = True) -> pd.DataFrame:
    df = df.copy()
    out = {"garch": [], "har": []}
    panels = {}
    for code in CODES:
        px = _series(frames["etf"], "code", code)
        drops = corporate_action_dates(px, _series(frames["idx"], "code", UNDERLYINGS[code][4]))
        r = daily_returns(px, drops)
        panels[code] = (pd.DatetimeIndex(px.index), r, har_frame(r))
    hs = []
    for u, t in zip(df["underlying"], df["origin"]):
        cal, r, har = panels[u]
        # trading days in (t, t + 30 days] on the exchange calendar (known in advance)
        h = int(cal.searchsorted(t + pd.Timedelta(days=STEP_DAYS), side="right") - cal.searchsorted(t, side="right"))
        hs.append(h)
        out["garch"].append(garch11_vol(r[r.index <= t].to_numpy(), h) if garch else np.nan)
        out["har"].append(har_vol(har, t))
    df["h"] = hs
    for k, v in out.items():
        df[k] = np.asarray(v, dtype=float)
    return df


def add_combinations(df: pd.DataFrame, models: tuple[str, ...] = ("garch", "har")) -> pd.DataFrame:
    """Per model, fitted on closed samples only: bias removed, a blend with IV - premium, and a + b*IV + c*model.

    The blends and regressions only define the common sample here (every
    forecast present), so this script scores exactly the published samples.
    """
    df = df.copy()
    y, iv, bench = df["rv_fwd"].to_numpy(), df["iv"].to_numpy(), df["iv_minus_premium"].to_numpy()
    for m in models:
        mv = df[m].to_numpy()
        db, mix, enc = (np.full(len(df), np.nan) for _ in range(3))
        for t, idx in df.groupby("origin").groups.items():
            ok = closed_before(df, t) & ~np.isnan(mv) & ~np.isnan(y)
            pos = df.index.get_indexer(idx)
            if ok.sum() >= OOS_MIN:
                db[pos] = mv[pos] - np.mean(mv[ok] - y[ok])
                X = np.column_stack([np.ones(ok.sum()), iv[ok], mv[ok]])
                a, b, c = np.linalg.lstsq(X, y[ok], rcond=None)[0]
                enc[pos] = a + b * iv[pos] + c * mv[pos]
            okb = ok & ~np.isnan(bench)
            if okb.sum() >= OOS_MIN:
                x = bench[okb] - mv[okb]
                w = float(((y[okb] - mv[okb]) @ x) / (x @ x))
                mix[pos] = w * bench[pos] + (1 - w) * mv[pos]
        df[m + "_debiased"] = np.maximum(db, FLOOR)
        df[m + "_blend"] = np.maximum(mix, FLOOR)
        df[m + "_regression"] = np.maximum(enc, FLOOR)
    return df


# ---------------------------------------------------------------------------
# Scores
# ---------------------------------------------------------------------------


def rmse_pts(y: pd.Series, f: pd.Series) -> float:
    return 100 * math.sqrt(((y - f) ** 2).mean())


def newey_west_t(x: np.ndarray, lags: int = 2) -> float:
    x = np.asarray(x, float)
    n = len(x)
    u = x - x.mean()
    s = u @ u / n
    for lag in range(1, lags + 1):
        s += 2 * (1 - lag / (lags + 1)) * (u[lag:] @ u[:-lag]) / n
    return float(x.mean() / math.sqrt(s / n))


def dm_t(df: pd.DataFrame, a: str, b: str) -> float:
    """Diebold-Mariano on squared error: per-date mean loss difference, Newey-West t. t > 0 = b better."""
    d = (df["rv_fwd"] - df[a]) ** 2 - (df["rv_fwd"] - df[b]) ** 2
    return newey_west_t(d.groupby(df["origin"]).mean().sort_index().to_numpy())


def calibration(df: pd.DataFrame) -> pd.DataFrame:
    """Log-normal range around IV - premium: log RV ~ N(log centre + m, s^2), m and s from past log errors."""
    df = df.copy()
    m = np.full(len(df), np.nan)
    s = np.full(len(df), np.nan)
    n_closed = np.zeros(len(df), dtype=int)
    iv, y = df["iv"].to_numpy(), df["rv_fwd"].to_numpy()
    for t, idx in df.groupby("origin").groups.items():
        ok = closed_before(df, t)
        if ok.sum() < OOS_MIN:
            continue
        pos = df.index.get_indexer(idx)
        prem = np.mean(iv[ok] - y[ok])
        err = np.log(y[ok] / np.maximum(iv[ok] - prem, FLOOR))
        m[pos], s[pos], n_closed[pos] = err.mean(), err.std(ddof=1), ok.sum()
    df["log_mean"], df["log_sd"], df["n_closed"] = m, s, n_closed
    z = np.abs(np.log(df["rv_fwd"]) - (np.log(df["iv_minus_premium"]) + df["log_mean"])) / df["log_sd"]
    for c in (0.5, 0.8, 0.9):
        df[f"in{int(c * 100)}"] = z <= statistics.NormalDist().inv_cdf((1 + c) / 2)
    return df


CLEAN_FROM = pd.Timestamp("2024-07-01")  # after the foundation models' training data ends


def compute(frames: dict, *, end: str | None = None, garch: bool = True, with_fm: bool = False) -> dict:
    """Every headline number. ``frames`` as returned by :func:`fetch`; ``end`` cuts the data at a date."""
    if end is not None:
        frames = {k: v[v["date"] <= pd.Timestamp(end)] for k, v in frames.items()}
    df = add_iv_forecasts(build_samples(frames))
    res = {
        "samples": len(df), "origins": df["origin"].nunique(),
        "rmse_trailing_rv": rmse_pts(df["rv_fwd"], df["fa"]),
        "rmse_qvix": rmse_pts(df["rv_fwd"], df["iv"]),
    }
    e = df.dropna(subset=["premium"])
    res["n_iv_minus_premium"] = len(e)
    res["rmse_iv_minus_premium_all"] = rmse_pts(e["rv_fwd"], e["iv_minus_premium"])

    df = add_combinations(add_models(df, frames, garch=garch)) if garch else add_models(df, frames, garch=False)
    # Without GARCH (--no-garch) the common sample is wider than the published one.
    every = ["iv_minus_premium", "fa", "iv", "a_b_iv", "har"] + (
        ["garch"] + [m + k for m in ("garch", "har") for k in ("_debiased", "_blend", "_regression")]
        if garch else [])
    common = df.dropna(subset=every)
    res["common_samples"], res["common_origins"] = len(common), common["origin"].nunique()
    bench = rmse_pts(common["rv_fwd"], common["iv_minus_premium"])
    table = {}
    for key in ("iv_minus_premium", "fa", "iv", "a_b_iv", "garch", "har", "garch_debiased", "har_debiased"):
        if key not in every:
            continue
        r = rmse_pts(common["rv_fwd"], common[key])
        table[key] = dict(rmse=r, ratio=r / bench,
                          dm_t=None if key == "iv_minus_premium" else dm_t(common, key, "iv_minus_premium"))
    res["common_table"] = table

    cal = calibration(df)
    for name, sub in (("clean", cal[(cal["origin"] >= CLEAN_FROM) & cal["log_sd"].notna()]),
                      ("full", cal[cal["n_closed"] >= 100])):
        res[f"coverage_{name}"] = dict(n=len(sub), origins=sub["origin"].nunique(),
                                       c50=sub["in50"].mean(), c80=sub["in80"].mean(), c90=sub["in90"].mean())
    last = cal.dropna(subset=["log_sd"]).iloc[-1]
    res["log_sd_now"], res["premium_now_pts"] = float(last["log_sd"]), 100 * float(last["premium"])

    if with_fm:
        res["timesfm"] = timesfm_clean(df, frames)
    res["frame"] = df
    return res


# ---------------------------------------------------------------------------
# Optional: TimesFM 2.5, zero-shot
# ---------------------------------------------------------------------------


def timesfm_clean(df: pd.DataFrame, frames: dict, max_context: int = 1024) -> dict:
    """TimesFM 2.5 on 252 x r^2 (last 1,024 returns), mean head averaged over the h days, from July 2024 on."""
    import timesfm  # pip install "timesfm[torch]>=3.0.2" "torch>=2.4"

    model = timesfm.TimesFM_2p5_200M_torch.from_pretrained("google/timesfm-2.5-200m-pytorch", torch_compile=False)
    model.compile(timesfm.ForecastConfig(
        max_context=max_context, max_horizon=128, normalize_inputs=True, use_continuous_quantile_head=True,
        force_flip_invariance=True, infer_is_positive=True, fix_quantile_crossing=True, per_core_batch_size=64))
    sub = df[(df["origin"] >= CLEAN_FROM) & df["premium"].notna()].copy()
    returns = {}
    for code in CODES:
        px = _series(frames["etf"], "code", code)
        returns[code] = daily_returns(px, corporate_action_dates(px, _series(frames["idx"], "code", UNDERLYINGS[code][4])))
    contexts = [252.0 * returns[u][returns[u].index <= t].to_numpy()[-max_context:] ** 2
                for u, t in zip(sub["underlying"], sub["origin"])]
    horizon = int(sub["h"].max())
    _, q = model.forecast(horizon=horizon, inputs=contexts)
    q = np.asarray(q)
    sub["timesfm"] = [max(math.sqrt(max(q[i, :h, 0].mean(), 0.0)), FLOOR) for i, h in enumerate(sub["h"])]
    rb = rmse_pts(sub["rv_fwd"], sub["iv_minus_premium"])
    rt = rmse_pts(sub["rv_fwd"], sub["timesfm"])
    return dict(n=len(sub), origins=sub["origin"].nunique(), rmse_iv_minus_premium=rb, rmse_timesfm=rt,
                ratio=rt / rb, dm_t=dm_t(sub, "timesfm", "iv_minus_premium"))


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

LABELS = {
    "iv_minus_premium": "IV - premium (the benchmark)",
    "fa": "trailing realised vol",
    "iv": "QVIX, raw",
    "a_b_iv": "a + b*IV (control)",
    "garch": "GARCH(1,1)",
    "har": "HAR-RV",
    "garch_debiased": "GARCH(1,1), bias removed",
    "har_debiased": "HAR-RV, bias removed",
}


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="Reproduce the vol-forecast study's headline numbers.")
    ap.add_argument("--cache-dir", help="save the download here as CSV and reuse it on the next run")
    ap.add_argument("--end", default="2026-09-30",
                    help="use data up to this date (default: the published run; 'none' = everything)")
    ap.add_argument("--no-garch", action="store_true", help="skip the GARCH fits (fast; no model table)")
    ap.add_argument("--with-fm", action="store_true", help="also run TimesFM 2.5 (needs timesfm + torch)")
    args = ap.parse_args(argv)

    end = None if args.end.lower() == "none" else args.end
    res = compute(fetch(args.cache_dir), end=end, garch=not args.no_garch, with_fm=args.with_fm)
    print(f"Samples: {res['samples']} on {res['origins']} dates "
          f"(IV - premium has {res['n_iv_minus_premium']}: it needs 24 closed months first)")
    print(f"RMSE, all samples (vol pts): trailing RV {res['rmse_trailing_rv']:.2f}, QVIX {res['rmse_qvix']:.2f}, "
          f"IV - premium {res['rmse_iv_minus_premium_all']:.2f}")
    print(f"\nCommon sample: {res['common_samples']} samples on {res['common_origins']} dates")
    print(f"{'forecast':32s} {'RMSE':>6s} {'ratio':>6s} {'DM t':>6s}")
    for key, row in res["common_table"].items():
        t = "" if row["dm_t"] is None else f"{row['dm_t']:+.2f}"
        print(f"{LABELS[key]:32s} {row['rmse']:6.2f} {row['ratio']:6.3f} {t:>6s}")
    print("\nCalibration of the log-normal range around IV - premium (share inside the 50 / 80 / 90 % range):")
    for name, label in (("clean", "from July 2024"), ("full", "full history")):
        c = res[f"coverage_{name}"]
        print(f"  {label:15s} {c['n']:4d} samples, {c['origins']:3d} dates: "
              f"{c['c50']:.2f} / {c['c80']:.2f} / {c['c90']:.2f}")
    print(f"  today's log sd {res['log_sd_now']:.3f}, premium {res['premium_now_pts']:.2f} pts (QVIX)")
    if "timesfm" in res:
        f = res["timesfm"]
        print(f"\nTimesFM 2.5 on squared returns, from July 2024 ({f['n']} samples, {f['origins']} dates): "
              f"RMSE {f['rmse_timesfm']:.2f} vs IV - premium {f['rmse_iv_minus_premium']:.2f}, "
              f"ratio {f['ratio']:.3f}, DM t {f['dm_t']:+.2f}")
    return 0


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