#!/usr/bin/env python3
"""Is past vol a better forecast than implied vol?  A self-contained backtest.

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, to deal in it without restriction, including the rights to
use, copy, modify, merge, publish, distribute, sublicense and/or sell copies,
subject to this notice being included in all copies. THE SOFTWARE IS PROVIDED
"AS IS", WITHOUT WARRANTY OF ANY KIND. For learning only; not investment advice.

What it does
------------
For six A-share ETFs with listed options, on a grid of dates at least 30
calendar days apart (2015-05 onward), it compares three forecasts of the
realised vol of the *next* 30 calendar days:

  (a) trailing RV        mean of the last 30, 60 and 90 days of realised vol
  (b) QVIX               the exchange-style 30-day implied-vol index
  (e) IV - premium       QVIX minus the average premium (IV - realised vol)
                         seen on samples whose 30-day window had already
                         closed (out of sample, expanding window)

and measures the premium IV - realised vol that an option seller collects.

Data (all public, through AKShare):
  ak.fund_etf_hist_sina(symbol)        ETF daily closes (Sina, unadjusted)
  ak.stock_zh_index_daily(symbol)      tracking-index closes (Sina); used only
                                       to spot ETF splits so they don't count
                                       as price moves
  ak.index_option_<x>_qvix()           QVIX daily closes (from 2015-02-09)

Run
---
  pip install -r requirements.txt
  python vrp_backtest.py                     # fetch from AKShare and print
  python vrp_backtest.py --cache-dir data    # save / reuse CSVs in ./data
  python vrp_backtest.py --until 2026-09-30  # cut the data at a date

Structure: fetch() -> dict of DataFrames, compute(frames) -> dict of numbers.
Vol is in decimals inside the code (0.20 = 20 %); the printout is in vol
points (1 pt = 0.01).
"""

from __future__ import annotations

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

import numpy as np
import pandas as pd

# ---------------------------------------------------------------------------
# Underlyings: one per QVIX series (QVIX code, Sina ETF symbol, tracking index)
# ---------------------------------------------------------------------------

UNDERLYINGS = {
    "510050": ("50etf", "sh510050", "sh000016"),   # 50ETF
    "510300": ("300etf", "sh510300", "sh000300"),  # 300ETF
    "510500": ("500etf", "sh510500", "sh000905"),  # 500ETF
    "159915": ("cyb", "sz159915", "sz399006"),     # ChiNext (创业板) ETF
    "588000": ("kcb", "sh588000", "sh000688"),     # STAR 50 (科创50) ETF
    "159901": ("100etf", "sz159901", "sz399330"),  # Shenzhen 100 ETF
}
QVIX_FUNCS = {
    "50etf": "index_option_50etf_qvix",
    "300etf": "index_option_300etf_qvix",
    "500etf": "index_option_500etf_qvix",
    "cyb": "index_option_cyb_qvix",
    "kcb": "index_option_kcb_qvix",
    "100etf": "index_option_100etf_qvix",
}

# ---------------------------------------------------------------------------
# Method constants
# ---------------------------------------------------------------------------

ANNUALISATION = 252
MIN_RETURNS = 10          # fewer valid returns than this -> no RV
MAX_GAP_DAYS = 5          # trailing RV skips returns spanning > 5 calendar days
TRAILING_WINDOWS = (30, 60, 90)
FORWARD_DAYS = 30
GRID_START = date(2015, 5, 11)   # first QVIX print (2015-02-09) + 90 days
LOOKBACK_SLICE = 95
OOS_MIN = 24              # samples needed before the out-of-sample fit starts
SIGMA_FLOOR = 0.01

# A return is a split / consolidation, not a price move, when the ETF moved
# > 10 % while its index stayed calm (< 3 %), the gap is > 5 %, and the move
# is not a +-10 % daily limit move.
CA_GAP, CA_MIN_MOVE, CA_INDEX_CALM, LIMIT_TOL = 0.05, 0.10, 0.03, 0.001


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


def fetch() -> dict[str, pd.DataFrame]:
    """Download QVIX, ETF closes and index closes from AKShare."""
    import akshare as ak

    qv, etf, idx = [], [], []
    for code, (series, etf_sym, idx_sym) in UNDERLYINGS.items():
        q = getattr(ak, QVIX_FUNCS[series])().rename(columns=str.lower)
        qv.append(pd.DataFrame({"series": series, "date": pd.to_datetime(q["date"]), "close": q["close"]}))
        e = ak.fund_etf_hist_sina(symbol=etf_sym)
        etf.append(pd.DataFrame({"code": code, "date": pd.to_datetime(e["date"]), "close": e["close"]}))
        i = ak.stock_zh_index_daily(symbol=idx_sym)
        idx.append(pd.DataFrame({"code": idx_sym[2:], "date": pd.to_datetime(i["date"]), "close": i["close"]}))
    return {"qvix": pd.concat(qv), "etf": pd.concat(etf), "index": pd.concat(idx)}


def load_or_fetch(cache_dir: Path | None) -> dict[str, pd.DataFrame]:
    if cache_dir is not None and all((cache_dir / f"{k}.csv").exists() for k in ("qvix", "etf", "index")):
        return {k: pd.read_csv(cache_dir / f"{k}.csv", dtype={"code": str, "series": str}, parse_dates=["date"])
                for k in ("qvix", "etf", "index")}
    frames = fetch()
    if cache_dir is not None:
        cache_dir.mkdir(parents=True, exist_ok=True)
        for k, df in frames.items():
            df.to_csv(cache_dir / f"{k}.csv", index=False)
    return frames


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, code: str) -> pd.Series:
    s = _series(frames["qvix"], "series", UNDERLYINGS[code][0])
    return s[s > 0] / 100.0  # QVIX is quoted in %, zero rows are gaps


def price_series(frames, code: str) -> pd.Series:
    return _series(frames["etf"], "code", code)


def index_series(frames, code: str) -> pd.Series:
    return _series(frames["index"], "code", UNDERLYINGS[code][2][2:])


def corporate_action_dates(frames, code: str) -> list[date]:
    df = pd.DataFrame({"px": price_series(frames, code), "ix": index_series(frames, code)}).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 (np.log(0.9), np.log(1.1)):
        at_limit |= ((r_px - lim).abs() < LIMIT_TOL).to_numpy()
    ca = ((r_px - r_ix).abs() > CA_GAP) & (r_ix.abs() < CA_INDEX_CALM) & (r_px.abs() > CA_MIN_MOVE) & ~at_limit
    return [ts.date() for ts in df.index[ca.fillna(False).to_numpy()]]


# ---------------------------------------------------------------------------
# Realised vol
# ---------------------------------------------------------------------------


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

    Log returns, population std (divide by n), x sqrt(252). ``skip_gaps``
    drops returns spanning more than 5 calendar days (long holidays), as the
    trailing forecast does; the forward target keeps every trading day.
    """
    start = as_of - timedelta(days=lookback_days)
    window = [(d, p) for d, p in points if start <= d <= as_of and p > 0]
    if len(window) < MIN_RETURNS + 1:
        return None
    drop = set(drop_dates)
    rets = [math.log(p1 / p0) for (d0, p0), (d1, p1) in zip(window, window[1:])
            if not (skip_gaps and (d1 - d0).days > MAX_GAP_DAYS) and d1 not in drop]
    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) * ANNUALISATION)


def trailing_rv(points, as_of: date, drop_dates=()) -> float | None:
    """Forecast (a): the mean of 30-, 60- and 90-day realised vol."""
    rvs = [realized_vol(points, as_of, w, skip_gaps=True, drop_dates=drop_dates) for w in TRAILING_WINDOWS]
    rvs = [v for v in rvs if v is not None]
    return max(sum(rvs) / len(rvs), SIGMA_FLOOR) if rvs else None


def forward_rv(points, origin: date, drop_dates=()) -> float | None:
    """Target: realised vol of the returns in (origin, origin + 30 days]."""
    return realized_vol(points, origin + timedelta(days=FORWARD_DAYS), FORWARD_DAYS,
                        skip_gaps=False, drop_dates=drop_dates)


# ---------------------------------------------------------------------------
# Samples
# ---------------------------------------------------------------------------


def origin_grid(calendar: list[date]) -> list[date]:
    """Trading days from GRID_START, each 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=FORWARD_DAYS))
    return out


def _points(s: pd.Series):
    pts = [(ts.date(), float(v)) for ts, v in s.items()]
    return [d for d, _ in pts], pts


def _slice(dates, pts, lo: date, hi: date):
    return pts[bisect.bisect_left(dates, lo): bisect.bisect_right(dates, hi)]


def build_samples(frames) -> pd.DataFrame:
    calendar = [ts.date() for ts in price_series(frames, "510050").index]
    grid = origin_grid(calendar)
    rows = []
    for code in UNDERLYINGS:
        qv = qvix_series(frames, code)
        q_map = {ts.date(): float(v) for ts, v in qv.items()}
        px_dates, px = _points(price_series(frames, code))
        drops = corporate_action_dates(frames, code)
        for t in grid:
            if t not in q_map or t + timedelta(days=FORWARD_DAYS) > px_dates[-1]:
                continue
            back = _slice(px_dates, px, t - timedelta(days=LOOKBACK_SLICE), t)
            fwd = _slice(px_dates, px, t, t + timedelta(days=FORWARD_DAYS))
            fa, target = trailing_rv(back, t, drops), forward_rv(fwd, t, drops)
            if fa is None or target is None:
                continue
            rows.append(dict(underlying=code, origin=pd.Timestamp(t), iv=q_map[t], fa=fa, rv_fwd=target))
    return pd.DataFrame(rows)


def add_iv_minus_premium(df: pd.DataFrame) -> pd.DataFrame:
    """Forecast (e): IV minus the mean premium of samples whose window has closed."""
    df = df.sort_values(["origin", "underlying"]).reset_index(drop=True)
    prem = np.full(len(df), np.nan)
    for t, grp in df.groupby("origin"):
        past = df[df["origin"] + pd.Timedelta(days=FORWARD_DAYS) <= t]
        if len(past) >= OOS_MIN:
            prem[grp.index] = float((past["iv"] - past["rv_fwd"]).mean())
    df["prem_oos"] = prem
    df["iv_minus_prem"] = (df["iv"] - df["prem_oos"]).clip(lower=SIGMA_FLOOR)
    return df


# ---------------------------------------------------------------------------
# Statistics
# ---------------------------------------------------------------------------


def error_stats(target: pd.Series, forecast: pd.Series) -> dict:
    """bias = mean(target - forecast) in vol pts (negative: forecast too high)."""
    m = target.notna() & forecast.notna()
    e, le = target[m] - forecast[m], np.log(target[m] / forecast[m])
    return dict(n=int(m.sum()), bias=100 * e.mean(), rmse=100 * math.sqrt((e ** 2).mean()),
                log_sd=float(le.std(ddof=1)))


def premium_stats(x: pd.Series) -> dict:
    x = x.dropna() * 100
    p5 = x.quantile(0.05)
    return dict(n=len(x), mean=x.mean(), median=x.median(), share_pos=(x > 0).mean(),
                sd=x.std(ddof=1), worst5_mean=x[x <= p5].mean())


def blend_weight(iv, rv, target) -> float:
    """w in target ~ w*IV + (1 - w)*trailing RV (no-intercept least squares)."""
    x, y = np.asarray(iv) - np.asarray(rv), np.asarray(target) - np.asarray(rv)
    return float(x @ y / (x @ x))


def compute(frames: dict[str, pd.DataFrame]) -> dict:
    """All headline numbers (vol points unless noted)."""
    df = add_iv_minus_premium(build_samples(frames))
    prem = df["iv"] - df["rv_fwd"]
    signal = df["iv"] - df["fa"]  # what a trailing-RV forecast calls the edge
    reverse = signal < 0           # trailing RV above IV: the "cheap IV" premise
    slope, intercept = np.polyfit(signal, df["rv_fwd"] - df["fa"], 1)
    log_err_e = np.log(df["rv_fwd"] / df["iv_minus_prem"]).dropna()
    return {
        "samples": len(df),
        "origins": int(df["origin"].nunique()),
        "first_origin": str(df["origin"].min().date()),
        "last_origin": str(df["origin"].max().date()),
        "forecasts": {
            "trailing_rv": error_stats(df["rv_fwd"], df["fa"]),
            "qvix": error_stats(df["rv_fwd"], df["iv"]),
            "iv_minus_premium_oos": error_stats(df["rv_fwd"], df["iv_minus_prem"]),
        },
        "premium_qvix": premium_stats(prem),
        "blend_w": blend_weight(df["iv"], df["fa"], df["rv_fwd"]),
        "blend_slope_with_intercept": float(slope),
        "blend_intercept_pts": 100 * float(intercept),
        "sell_signal": dict(n=int((signal > 0).sum()), hit_rate=float((prem[signal > 0] > 0).mean()),
                            edge=100 * float(prem[signal > 0].mean())),
        "cheap_iv_premise": dict(n=int(reverse.sum()), hit_rate=float((prem[reverse] < 0).mean()),
                                 buyer_edge=100 * float(-prem[reverse].mean())),
        "error_spread_e": dict(log_sd=float(log_err_e.std(ddof=1)), p5=float(log_err_e.quantile(0.05)),
                               p95=float(log_err_e.quantile(0.95))),
    }


def cut(frames: dict[str, pd.DataFrame], until: str | None) -> dict[str, pd.DataFrame]:
    if until is None:
        return frames
    return {k: v[v["date"] <= pd.Timestamp(until)] for k, v in frames.items()}


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description="VRP forecast backtest (public data)")
    ap.add_argument("--cache-dir", type=Path, help="save / reuse the downloaded CSVs here")
    ap.add_argument("--until", help="ignore data after this date (YYYY-MM-DD)")
    args = ap.parse_args(argv)
    r = compute(cut(load_or_fetch(args.cache_dir), args.until))
    print(f"{r['samples']} samples on {r['origins']} origin dates, {r['first_origin']} .. {r['last_origin']}\n")
    print(f"{'forecast':<24}{'n':>5}{'bias':>8}{'RMSE':>8}{'log sd':>8}")
    for name, s in r["forecasts"].items():
        print(f"{name:<24}{s['n']:>5}{s['bias']:>+8.2f}{s['rmse']:>8.2f}{s['log_sd']:>8.3f}")
    p = r["premium_qvix"]
    print(f"\npremium QVIX - realised: mean {p['mean']:+.2f}, median {p['median']:+.2f}, "
          f"positive {p['share_pos']:.0%} of months, sd {p['sd']:.2f}, worst 5% mean {p['worst5_mean']:+.2f}")
    print(f"blend weight on IV w = {r['blend_w']:.3f}; with an intercept: "
          f"next RV - trailing RV = {r['blend_slope_with_intercept']:.3f} x (IV - trailing RV) "
          f"{r['blend_intercept_pts']:+.2f} pts")
    s, c = r["sell_signal"], r["cheap_iv_premise"]
    print(f"sell when IV > trailing RV: {s['n']} months, IV above realised {s['hit_rate']:.0%}, premium {s['edge']:+.2f}")
    print(f"'cheap IV' (trailing RV > IV): {c['n']} months, realised above IV {c['hit_rate']:.0%}, "
          f"buyer's edge {c['buyer_edge']:+.2f}")
    e = r["error_spread_e"]
    print(f"IV - premium log error: sd {e['log_sd']:.3f}, p5 {e['p5']:+.3f}, p95 {e['p95']:+.3f}")
    return 0


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